Update main.py

main.py

@@ -6,20 +6,20 @@ from fastapi import FastAPI, Query
 from contextlib import asynccontextmanager
 
 DB_PATH = "insta_hitech.db"
-# Aapka dataset direct download link (Token ke saath)
+# Dataset direct download link
 DB_URL = "https://huggingface.co/datasets/Watchhrr/Insta17m/resolve/main/insta_hitech.db"
 TOKEN = os.getenv("MY_HF_TOKEN")
 
 def download_db():
     if not os.path.exists(DB_PATH):
-        print("📥 Database missing. Downloading from Dataset...")
-        headers = {"Authorization": f"Bearer {TOKEN}"}
+        print("📥 Downloading DB from Dataset...")
+        headers = {"Authorization": f"Bearer {TOKEN}"} if TOKEN else {}
         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=1024*1024):
                     f.write(chunk)
-        print("✅ Download Complete!")
+        print("✅ DB Downloaded!")
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -28,38 +28,36 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(lifespan=lifespan)
 
-@app.get("/")
-def home():
-    return {"message": "Insta Search API is Live!", "endpoint": "/search?q=YOUR_QUERY"}
+# 1. Total Count Check karne ke liye: /count
+@app.get("/count")
+def get_count():
+    if not os.path.exists(DB_PATH): return {"error": "DB not ready"}
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    cursor.execute("SELECT COUNT(*) FROM users")
+    total = cursor.fetchone()[0]
+    conn.close()
+    return {"total_records": total}
 
+# 2. Main Search: /search?q=XYZ
 @app.get("/search")
 async def search(q: str = Query(None)):
-    if not q:
-        return {"error": "Bhai, search query 'q' miss hai. Example: /search?q=username"}
-    
-    if not os.path.exists(DB_PATH):
-        return {"error": "Database file abhi download ho rahi hai, thoda sabar rakhein."}
+    if not q: return {"error": "Bhai, 'q' parameter missing hai!"}
+    if not os.path.exists(DB_PATH): return {"error": "DB downloading..."}
 
     conn = sqlite3.connect(DB_PATH)
-    conn.row_factory = sqlite3.Row # JSON format output ke liye
+    conn.row_factory = sqlite3.Row
     cursor = conn.cursor()
     
     try:
-        # Optimized query jo indexes use karegi
-        # 'u' (username), 'id' (user_id), 't' (phone) teeno check honge
-        sql_query = "SELECT * FROM users WHERE u = ? OR id = ? OR t = ? LIMIT 10"
-        cursor.execute(sql_query, (q, q, q))
+        # Search query jo Username(u), ID(id), aur Phone(t) teeno check karega
+        cursor.execute("SELECT * FROM users WHERE u=? OR id=? OR t=? LIMIT 10", (q, q, q))
         rows = cursor.fetchall()
         
         if not rows:
-            return {"status": "error", "message": "No Record Found in Database"}
+            return {"status": "error", "message": "No Record Found"}
         
-        # Result ko list of dict mein badalna
-        results = [dict(row) for row in rows]
-        return {"status": "success", "count": len(results), "data": results}
-    
-    except Exception as e:
-        return {"status": "error", "message": str(e)}
+        return {"status": "success", "results": [dict(r) for r in rows]}
     finally:
         conn.close()