Sheeturt's picture
Update app.py
d604070 verified
Raw
History Blame Contribute Delete
2 kB
from fastapi import FastAPI
import sqlite3
import os
import shutil
from huggingface_hub import hf_hub_download
app = FastAPI()
DATASET_REPO = "sheeturt/boat-master-data"
DB_NAME = "boat_master.db"
def join_db():
# Agar DB pehle se join ho chuka hai toh skip karo
if os.path.exists(DB_NAME) and os.path.getsize(DB_NAME) > 2000000000:
return
print("🛠 Downloading and Joining DB parts...")
try:
with open(DB_NAME, 'wb') as master:
for i in range(1, 5): # Check all 4 parts
part_path = hf_hub_download(repo_id=DATASET_REPO, filename=f"boat_master.db.part{i}", repo_type="dataset")
with open(part_path, 'rb') as part:
shutil.copyfileobj(part, master) # Efficient copying
print("✅ Join Complete!")
except Exception as e:
print(f"❌ Join Error: {e}")
@app.on_event("startup")
async def startup_event():
join_db()
@app.get("/search")
def search(num: str = None, email: str = None):
try:
if not os.path.exists(DB_NAME):
return {"error": "Database files are being downloaded. Try in 5 minutes."}
conn = sqlite3.connect(DB_NAME)
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
if num:
# Query with LIMIT to prevent memory overflow
cursor.execute("SELECT phone, email, raw_info FROM boat_data WHERE phone LIKE ? LIMIT 50", (f"%{num}%",))
elif email:
cursor.execute("SELECT phone, email, raw_info FROM boat_data WHERE email = ? LIMIT 50", (email,))
else:
return {"error": "Please provide 'num' or 'email' parameter."}
res = [dict(row) for row in cursor.fetchall()]
conn.close()
return {"count": len(res), "results": res}
except Exception as e:
return {"error": str(e)}
@app.get("/")
def home():
return {"status": "Running", "message": "Use /search?num=123 or /search?email=abc"}