| 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 |
| import pyarrow.dataset as ds |
| import pyarrow.compute as pc |
| import pyarrow.parquet as pq |
|
|
| logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") |
| logger = logging.getLogger(__name__) |
|
|
| REPO_ID = "sauravsingh2111/Tgdata" |
| 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): |
| for n in names: |
| if n.endswith(".parquet"): |
| files.append(os.path.join(root, n)) |
| return sorted(files) |
|
|
| class UserIdIndex: |
| def __init__(self): |
| self.entries = [] |
| self.keys = [] |
|
|
| def build(self, arrow_dataset): |
| logger.info("Building user_id index from parquet metadata...") |
| t0 = time.time() |
| schema = arrow_dataset.schema |
| uid_idx = schema.get_field_index("user_id") |
| if uid_idx < 0: |
| logger.error("user_id column not found!") |
| return |
| for frag in arrow_dataset.get_fragments(): |
| meta = frag.metadata |
| path = frag.path |
| for rg_idx in range(meta.num_row_groups): |
| col = meta.row_group(rg_idx).column(uid_idx) |
| s = col.statistics |
| if s and s.min is not None and s.max is not None: |
| self.entries.append((int(s.min), int(s.max), path, rg_idx)) |
| self.entries.sort(key=lambda x: x[0]) |
| self.keys = [e[0] for e in self.entries] |
| elapsed = time.time() - t0 |
| logger.info(f"Index built: {len(self.entries)} row groups in {elapsed:.1f}s") |
|
|
| def find(self, user_id): |
| idx = bisect.bisect_right(self.keys, user_id) - 1 |
| 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, is_ready, init_error |
|
|
| try: |
| os.makedirs(CACHE_DIR, exist_ok=True) |
| files = find_parquet(CACHE_DIR) |
|
|
| 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=["*"]) |
|
|
| 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 |
| t0 = time.time() |
|
|
| result = user_id_index.find(user_id) |
| if result: |
| path, rg_idx = result |
| try: |
| pf = pq.ParquetFile(path) |
| tbl = pf.read_row_group(rg_idx, columns=ALL_COLS) |
| tbl = tbl.filter(pc.equal(pc.field("user_id"), user_id)) |
| if len(tbl): |
| 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 |
| return (tbl.to_pylist()[0], elapsed) if len(tbl) else (None, elapsed) |
|
|
| def search_users(params: dict, limit: int = 10): |
| stats["queries"] += 1 |
| t0 = time.time() |
|
|
| conditions = [] |
| for key in ("username", "phone", "email", "first_name", "last_name"): |
| val = params.get(key) |
| if val: |
| conditions.append(pc.equal(pc.field(key), val)) |
|
|
| if not conditions: |
| return [], 0, time.time() - t0 |
|
|
| combined = conditions[0] |
| 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"]) |
| elapsed = time.time() - t0 |
| count = len(tbl) |
| if not count: |
| return [], 0, elapsed |
| return tbl.to_pylist()[:limit], count, elapsed |
|
|
| @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", |
| "/health": "System health & stats" |
| } |
| } |
|
|
| @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), |
| phone: str = Query(None), |
| email: str = Query(None), |
| first_name: str = Query(None), |
| 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") |
|
|
| rows, total, elapsed = search_users(params, limit) |
| if not rows: |
| 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": "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"], |
| "row_groups": stats["total_row_groups"], |
| "rows": stats["total_rows"], |
| "queries_served": stats["queries"], |
| "uptime_s": uptime |
| } |
|
|
