Update main.py

main.py

@@ -7,65 +7,51 @@ from contextlib import asynccontextmanager
 
 DB_PATH = "insta_hitech.duckdb"
 DB_URL = "https://huggingface.co/datasets/Watchhrr/Insta17m/resolve/main/insta_hitech.duckdb"
+TOKEN = os.getenv("MY_HF_TOKEN")
 
-# Yahan humne secret variable ka naam use kiya hai jo aapne Settings mein dala tha
-TOKEN = os.getenv("MY_HF_TOKEN") 
-
-# 1. Download Function
 def download_db():
     if not os.path.exists(DB_PATH):
-        print("📥 Database missing. Downloading from Dataset...")
-        if not TOKEN:
-            print("❌ Error: MY_HF_TOKEN not found in Secrets!")
-            return
-            
-        headers = {"Authorization": f"Bearer {TOKEN}"}
-        try:
-            with requests.get(DB_URL, headers=headers, stream=True) as r:
-                r.raise_for_status()
-                with open(DB_PATH, 'wb') as f:
-                    for chunk in r.iter_content(chunk_size=8192):
+        print("📥 Database missing. Starting download...")
+        headers = {"Authorization": f"Bearer {TOKEN}"} if TOKEN else {}
+        
+        response = requests.get(DB_URL, headers=headers, stream=True)
+        total_size = int(response.headers.get('content-length', 0))
+        
+        if response.status_code == 200:
+            with open(DB_PATH, 'wb') as f:
+                downloaded = 0
+                for chunk in response.iter_content(chunk_size=1024*1024): # 1MB chunks
+                    if chunk:
                         f.write(chunk)
+                        downloaded += len(chunk)
+                        # Har 50MB par log dikhayega
+                        if (downloaded // (50*1024*1024)) > ((downloaded - len(chunk)) // (50*1024*1024)):
+                            done = int(50 * downloaded / total_size)
+                            print(f"⏳ Progress: [{'=' * done}{' ' * (50-done)}] {downloaded/(1024*1024):.1f}MB / {total_size/(1024*1024):.1f}MB")
             print("✅ Download Complete!")
-        except Exception as e:
-            print(f"❌ Download Failed: {e}")
+        else:
+            print(f"❌ Download Failed! Status: {response.status_code}. Check your Token/Secret.")
     else:
-        print(f"✅ Database found! Size: {os.path.getsize(DB_PATH)} bytes")
+        print(f"✅ DB found. Size: {os.path.getsize(DB_PATH)/(1024*1024):.1f} MB")
 
-# 2. Lifespan Handler
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     download_db()
     yield
-    print("Shutting down...")
 
 app = FastAPI(lifespan=lifespan)
 
 @app.get("/search")
 async def search(q: str = Query(None)):
-    if not q:
-        return {"error": "Query parameter 'q' is missing"}
+    if not q: return {"error": "Query 'q' missing"}
+    if not os.path.exists(DB_PATH): return {"error": "DB still downloading..."}
     
-    if not os.path.exists(DB_PATH):
-        return {"error": "Database file not ready yet. Please wait a few minutes."}
-
     conn = duckdb.connect(DB_PATH, read_only=True)
     try:
-        # Smart search: username (u), id, and phone (t)
-        query = "SELECT * FROM users WHERE u = ? OR id = ? OR t = ?"
-        res = conn.execute(query, [q, q, q]).df()
-        
-        if res.empty:
-            return {"status": "error", "message": "No Record Found"}
-        
-        data = res.rename(columns={
-            't':'phone', 'u':'username', 'e':'email', 
-            'n':'name', 'id':'user_id', 'a':'extra'
-        }).to_dict(orient='records')
-        
+        res = conn.execute("SELECT * FROM users WHERE u = ? OR id = ? OR t = ?", [q, q, q]).df()
+        if res.empty: return {"status": "error", "message": "No Record Found"}
+        data = res.rename(columns={'t':'phone', 'u':'username', 'e':'email', 'n':'name', 'id':'user_id'}).to_dict(orient='records')
         return {"status": "success", "data": data}
-    except Exception as e:
-        return {"status": "error", "message": str(e)}
     finally:
         conn.close()