Upload app.py with huggingface_hub

app.py

@@ -1,7 +1,9 @@
 import os
 import time
 import logging
+import asyncio
 import bisect
+from contextlib import asynccontextmanager
 from fastapi import FastAPI, HTTPException, Query
 from fastapi.middleware.cors import CORSMiddleware
 from huggingface_hub import snapshot_download
@@ -17,8 +19,12 @@ CACHE_DIR = "/data/tgdb_cache"
 
 dataset = None
 user_id_index = None
+is_ready = False
+init_error = None
 stats = {"startup_time": 0, "total_files": 0, "total_row_groups": 0, "total_rows": 0, "queries": 0}
 
+ALL_COLS = ["user_id", "username", "first_name", "last_name", "phone", "email", "status", "linked_id", "linked_name", "linked_handle"]
+
 def find_parquet(base):
     files = []
     for root, _, names in os.walk(base):
@@ -27,7 +33,6 @@ def find_parquet(base):
                 files.append(os.path.join(root, n))
     return sorted(files)
 
-
 class UserIdIndex:
     def __init__(self):
         self.entries = []
@@ -39,7 +44,7 @@ class UserIdIndex:
         schema = arrow_dataset.schema
         uid_idx = schema.get_field_index("user_id")
         if uid_idx < 0:
-            logger.error("user_id column not found in schema!")
+            logger.error("user_id column not found!")
             return
         for frag in arrow_dataset.get_fragments():
             meta = frag.metadata
@@ -56,49 +61,71 @@ class UserIdIndex:
 
     def find(self, user_id):
         idx = bisect.bisect_right(self.keys, user_id) - 1
-        if idx >= 0:
+        if idx >= 0 and idx < len(self.entries):
             mn, mx, path, rg = self.entries[idx]
             if mn <= user_id <= mx:
                 return path, rg
         return None
 
-
 def init_dataset():
-    global dataset, user_id_index
-
-    os.makedirs(CACHE_DIR, exist_ok=True)
-    files = find_parquet(CACHE_DIR)
-
-    if not files:
-        logger.info("Downloading dataset...")
-        snapshot_download(
-            repo_id=REPO_ID,
-            repo_type="dataset",
-            local_dir=CACHE_DIR,
-            local_dir_use_symlinks=False,
-        )
-        files = find_parquet(CACHE_DIR)
-        logger.info(f"Downloaded {len(files)} parquet files")
-
-    logger.info(f"Creating dataset from {len(files)} files...")
-    dataset = ds.dataset(files, format="parquet")
+    global dataset, user_id_index, is_ready, init_error
 
-    rg_count = sum(frag.metadata.num_row_groups for frag in dataset.get_fragments())
-    row_count = sum(frag.count_rows() for frag in dataset.get_fragments())
-
-    stats["total_files"] = len(files)
-    stats["total_row_groups"] = rg_count
-    stats["total_rows"] = row_count
-
-    logger.info(f"Dataset: {len(files)} files, {rg_count} row groups, {row_count:,} rows")
-
-    user_id_index = UserIdIndex()
-    user_id_index.build(dataset)
+    try:
+        os.makedirs(CACHE_DIR, exist_ok=True)
+        files = find_parquet(CACHE_DIR)
 
-    stats["startup_time"] = time.time()
+        if not files:
+            logger.info("Downloading dataset (~8GB)...")
+            t0 = time.time()
+            snapshot_download(
+                repo_id=REPO_ID,
+                repo_type="dataset",
+                local_dir=CACHE_DIR,
+                local_dir_use_symlinks=False,
+            )
+            logger.info(f"Download completed in {time.time()-t0:.1f}s")
+            files = find_parquet(CACHE_DIR)
+
+        logger.info(f"Creating dataset from {len(files)} files...")
+        d = ds.dataset(files, format="parquet")
+
+        rg_count = sum(f.metadata.num_row_groups for f in d.get_fragments())
+        row_count = sum(f.count_rows() for f in d.get_fragments())
+
+        idx = UserIdIndex()
+        idx.build(d)
+
+        stats["total_files"] = len(files)
+        stats["total_row_groups"] = rg_count
+        stats["total_rows"] = row_count
+        stats["startup_time"] = time.time()
+
+        dataset = d
+        user_id_index = idx
+        is_ready = True
+
+        logger.info(f"Ready: {len(files)} files, {rg_count} row groups, {row_count:,} rows")
+    except Exception as e:
+        init_error = str(e)
+        logger.error(f"Init failed: {e}")
+
+async def background_init():
+    loop = asyncio.get_event_loop()
+    await loop.run_in_executor(None, init_dataset)
+
+@asynccontextmanager
+async def lifespan(app):
+    asyncio.create_task(background_init())
+    yield
+
+app = FastAPI(title="Telegram DB API", description="859M Telegram users searchable by ID", version="1.0.0", lifespan=lifespan)
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
 
-
-ALL_COLS = ["user_id", "username", "first_name", "last_name", "phone", "email", "status", "linked_id", "linked_name", "linked_handle"]
+def require_ready():
+    if not is_ready:
+        if init_error:
+            raise HTTPException(status_code=503, detail=f"Init error: {init_error}")
+        raise HTTPException(status_code=503, detail="Dataset loading, please retry in a moment")
 
 def query_user(user_id: int):
     stats["queries"] += 1
@@ -112,17 +139,13 @@ def query_user(user_id: int):
             tbl = pf.read_row_group(rg_idx, columns=ALL_COLS)
             tbl = tbl.filter(pc.equal(pc.field("user_id"), user_id))
             if len(tbl):
-                elapsed = time.time() - t0
-                return tbl.to_pylist()[0], elapsed
+                return tbl.to_pylist()[0], time.time() - t0
         except Exception as e:
             logger.warning(f"Index lookup failed for {user_id}: {e}")
 
     tbl = dataset.to_table(filter=pc.equal(pc.field("user_id"), user_id), columns=ALL_COLS)
     elapsed = time.time() - t0
-    if len(tbl) == 0:
-        return None, elapsed
-    return tbl.to_pylist()[0], elapsed
-
+    return (tbl.to_pylist()[0], elapsed) if len(tbl) else (None, elapsed)
 
 def search_users(params: dict, limit: int = 10):
     stats["queries"] += 1
@@ -141,31 +164,19 @@ def search_users(params: dict, limit: int = 10):
     for c in conditions[1:]:
         combined = combined & c
 
-    tbl = dataset.to_table(
-        filter=combined,
-        columns=["user_id", "username", "first_name", "last_name", "phone", "email", "status"]
-    )
+    tbl = dataset.to_table(filter=combined, columns=["user_id", "username", "first_name", "last_name", "phone", "email", "status"])
     elapsed = time.time() - t0
     count = len(tbl)
-    if count == 0:
+    if not count:
         return [], 0, elapsed
     return tbl.to_pylist()[:limit], count, elapsed
 
-
-app = FastAPI(title="Telegram DB API", description="859M Telegram users searchable by ID", version="1.0.0")
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
-
-
-@app.on_event("startup")
-async def startup():
-    init_dataset()
-
-
 @app.get("/")
 async def root():
     return {
         "name": "Telegram DB API",
         "dataset": REPO_ID,
+        "ready": is_ready,
         "endpoints": {
             "/user/{user_id}": "Get full details by Telegram user ID",
             "/search": "Search by username, phone, email, first_name, last_name",
@@ -173,15 +184,14 @@ async def root():
         }
     }
 
-
 @app.get("/user/{user_id}")
 async def get_user(user_id: int):
+    require_ready()
     user, elapsed = query_user(user_id)
     if user is None:
         raise HTTPException(status_code=404, detail=f"User {user_id} not found")
     return {"found": True, "query_time_ms": round(elapsed * 1000, 2), "user": user}
 
-
 @app.get("/search")
 async def search(
     username: str = Query(None),
@@ -191,6 +201,7 @@ async def search(
     last_name: str = Query(None),
     limit: int = Query(10, ge=1, le=100),
 ):
+    require_ready()
     params = {k: v for k, v in {"username": username, "phone": phone, "email": email, "first_name": first_name, "last_name": last_name}.items() if v}
     if not params:
         raise HTTPException(status_code=400, detail="Provide at least one search parameter")
@@ -200,12 +211,13 @@ async def search(
         raise HTTPException(status_code=404, detail="No users found")
     return {"found": True, "count": total, "returned": len(rows), "query_time_ms": round(elapsed * 1000, 2), "users": rows}
 
-
 @app.get("/health")
 async def health():
     uptime = round(time.time() - stats["startup_time"], 1) if stats["startup_time"] else 0
     return {
-        "status": "ok",
+        "status": "loading" if not is_ready else ("error" if init_error else "ok"),
+        "ready": is_ready,
+        "error": init_error,
         "dataset": REPO_ID,
         "cached": os.path.exists(CACHE_DIR),
         "files": stats["total_files"],