"""FastAPI backend over DuckDB - ICMR + HITEK combined dataset (2.5B rows, 11 cols). Index-aware routing: - idx_phone.parquet (sorted by phoneNumber) -> fast exact phone/other lookups - idx_aadhar.parquet (sorted by aadharNumber) -> fast exact aadhar lookups - idx_name.parquet (sorted by name) -> fast name prefix/exact - Remote Parquet loading via Hugging Face Datasets with Token Authentication Dedup: max 2 rows per person. source column included (icmr / inddata "connected docs"). """ import asyncio import json import os import threading from concurrent.futures import ThreadPoolExecutor from typing import Any import duckdb from fastapi import FastAPI, HTTPException, Query, Response from pydantic import BaseModel BASE = os.path.dirname(os.path.abspath(__file__)) # ── Updated Database, Token & Remote Dataset URL ─────────────────────────── HF_TOKEN = os.environ.get("HF_TOKEN", "") HF_INDEX_BASE = os.environ.get( "ICMR_HF_INDEX_BASE", "https://huggingface.co/datasets/HiTeckGroup/HiTeckgroop/resolve/main", ).rstrip("/") INDEX_SOURCE = os.environ.get("ICMR_INDEX_SOURCE", "remote").lower() PARALLELISM = int(os.environ.get("ICMR_PARALLEL", "2")) THREADS_PER_CONN = int(os.environ.get("ICMR_THREADS_PER_CONN", "2")) DUPLICATE_CAP = 2 # max 2 copies of the same person in results SEARCH_FIELDS = [ "name", "fathersName", "phoneNumber", "aadharNumber", "otherNumber", "address", "district", "pincode", "state", "town", "source", ] NUMBER_FIELDS = ["phoneNumber", "aadharNumber", "otherNumber"] # Remote Parquet Files List REMOTE_INDEXES = { "phone": [f"{HF_INDEX_BASE}/idx_phone.{i}.parquet" for i in range(7)], "aadhar": [f"{HF_INDEX_BASE}/idx_aadhar.{i}.parquet" for i in range(7)], } # 🔑 API KEY (secret से load) API_KEY = os.environ.get("API_KEY_1", "") VALID_KEYS = {} if API_KEY: VALID_KEYS[API_KEY] = {} app = FastAPI(title="ICMR + HITEK Search API", description="DuckDB-backed search over 2.5B records (11 cols, dedup max 2)") # ---- DuckDB connection pool (one connection per worker thread) ---- _conns: list[duckdb.DuckDBPyConnection] = [] _conns_lock = threading.Lock() _thread_local = threading.local() pool = ThreadPoolExecutor(max_workers=PARALLELISM, thread_name_prefix="duck") def _idx_ready(kind: str) -> bool: return kind in REMOTE_INDEXES def _new_conn() -> duckdb.DuckDBPyConnection: con = duckdb.connect() # in-memory con.execute("INSTALL parquet; LOAD parquet;") con.execute("INSTALL httpfs; LOAD httpfs;") # ── PRIVATE DATASET AUTHENTICATION ───────────────────────────────────── if HF_TOKEN: con.execute(f"SET http_headers = {{'Authorization': 'Bearer {HF_TOKEN}'}};") # 🔥🔥🔥 SPEED OPTIMIZATIONS (नया) 🔥🔥🔥 con.execute("SET enable_http_metadata_cache = true;") # metadata cache con.execute("SET http_keep_alive = true;") # connection reuse con.execute("SET http_timeout = 600000;") # 10 मिनट timeout # Create sorted index views from remote Hugging Face dataset parts for kind, urls in REMOTE_INDEXES.items(): view = f"people_{kind}" lst = ", ".join(f"'{u}'" for u in urls) con.execute(f"CREATE OR REPLACE VIEW {view} AS SELECT * FROM read_parquet([{lst}])") con.execute(f"SET threads = {THREADS_PER_CONN}") return con def _thread_id() -> int: tid = getattr(_thread_local, "id", None) if tid is None: with _conns_lock: tid = len(_conns) _thread_local.id = tid return tid def _get_conn() -> duckdb.DuckDBPyConnection: ident = _thread_id() with _conns_lock: while len(_conns) <= ident: _conns.append(_new_conn()) return _conns[ident] # ---- Dedup: max 2 copies per person ---- def _person_key(row: dict[str, Any]) -> tuple[str, ...]: ph = (row.get("phoneNumber") or "").strip() ad = (row.get("aadharNumber") or "").strip() if ph or ad: return (ph, ad) return (row.get("name") or "").strip(), (row.get("fathersName") or "").strip() def _cap_duplicates(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: """Keep at most DUPLICATE_CAP rows per person (same phone+aadhar identity).""" seen: dict[tuple[str, ...], int] = {} out: list[dict[str, Any]] = [] for r in rows: k = _person_key(r) n = seen.get(k, 0) if n < DUPLICATE_CAP: seen[k] = n + 1 out.append(r) return out def _source_clause(source: str | None) -> str: if source and source in ("icmr", "hitek", "inddata"): src = "hitek" if source == "hitek" else source return f" AND source = '{src}'" return "" def _run_field_search(field: str, value: str, mode: str, limit: int, source: str | None = None) -> dict[str, Any]: if field not in SEARCH_FIELDS: raise ValueError(f"Unknown field: {field}") v = value.replace("'", "''") # Use remote sorted index view when available if mode == "exact": if field == "phoneNumber" and _idx_ready("phone"): view = "people_phone" elif field == "aadharNumber" and _idx_ready("aadhar"): view = "people_aadhar" else: return {"field": field, "value": value, "mode": mode, "count": 0, "results": []} sql = (f"SELECT * FROM {view} WHERE {field} = '{v}'" f"{_source_clause(source)} LIMIT {limit * DUPLICATE_CAP + 20}") elif mode == "contains": v2 = v.replace("%", r"\%").replace("_", r"\_") sql = (f"SELECT * FROM people_phone WHERE {field} ILIKE '%{v2}%' ESCAPE '\\'" f"{_source_clause(source)} LIMIT {limit * DUPLICATE_CAP + 20}") else: raise ValueError(f"Unknown mode: {mode}") con = _get_conn() rows = con.execute(sql).fetchall() cols = [d[0] for d in con.description] results = _cap_duplicates([dict(zip(cols, r)) for r in rows])[:limit] return {"field": field, "value": value, "mode": mode, "count": len(results), "results": results} # 🔑 KEY CHECK FUNCTION def _check_key(key: str | None) -> tuple[bool, dict | None]: if not key or key not in VALID_KEYS: return False, {"message": "Provide correct key", "contact": "@dmforbans"} return True, VALID_KEYS[key] async def _parallel(queries: list[tuple[str, str, str, int]]) -> list[dict[str, Any]]: loop = asyncio.get_running_loop() async def one(t: tuple[str, str, str, int]) -> dict[str, Any]: field, value, mode, limit = t return await loop.run_in_executor(pool, _run_field_search, field, value, mode, limit) return await asyncio.gather(*[one(t) for t in queries]) async def _unified_search(q: str, limit: int, source: str | None = None) -> dict[str, Any]: """Smart unified search: uses remote sorted indexes per field, merges, dedups.""" q = q.strip() is_num = q.isdigit() and len(q) >= 8 loop = asyncio.get_running_loop() if is_num: all_rows: list[dict[str, Any]] = [] searched: list[str] = [] if _idx_ready("phone"): r = await loop.run_in_executor(pool, _run_field_search, "phoneNumber", q, "exact", limit, source) all_rows.extend(r["results"]) searched.append("phoneNumber") if not all_rows and _idx_ready("aadhar"): r = await loop.run_in_executor(pool, _run_field_search, "aadharNumber", q, "exact", limit, source) all_rows.extend(r["results"]) searched.append("aadharNumber") all_rows = _cap_duplicates(all_rows)[:limit] return {"query": q, "searched_fields": searched or list(NUMBER_FIELDS), "dedup": DUPLICATE_CAP, "count": len(all_rows), "results": all_rows} else: return {"query": q, "searched_fields": [], "dedup": DUPLICATE_CAP, "count": 0, "results": []} def _pretty(data: dict[str, Any], pretty: bool) -> Response: if pretty: return Response(content=json.dumps(data, indent=2, ensure_ascii=False), media_type="application/json") return Response(content=json.dumps(data, ensure_ascii=False), media_type="application/json") # ---- API Routes ---- class BatchRequest(BaseModel): queries: list[dict[str, Any]] limit: int = 10 key: str | None = None # 🔑 KEY @app.get("/") def root(): return { "app": "ICMR + HITEK Search API", "status": "protected", "contact": "@dmforbans", } @app.get("/health") def health(): return { "status": "ok", "protected": True, } @app.get("/search") async def search( q: str = Query(..., description="Search term - phone number ya identification number"), key: str = Query(..., description="API Key"), # 🔑 KEY REQUIRED field: str | None = Query(None, description=f"Restrict to one field: {SEARCH_FIELDS}"), mode: str = Query("exact", pattern="^(exact|contains)$"), limit: int = Query(10, ge=1, le=1000), source: str | None = Query(None, pattern="^(icmr|hitek|inddata)$", description="Filter by source: icmr, hitek (=inddata), inddata"), pretty: bool = Query(True, description="Pretty-print JSON"), ): """Unified search across remote HF Parquet files.""" # 🔑 KEY CHECK valid, err = _check_key(key) if not valid: return Response( content=json.dumps(err, ensure_ascii=False), status_code=403, media_type="application/json" ) if field: try: data = _run_field_search(field, q, mode, limit, source) except ValueError as e: raise HTTPException(400, str(e)) data = {"query": q, "field": field, "mode": mode, "dedup": DUPLICATE_CAP, "count": data["count"], "results": data["results"]} else: data = await _unified_search(q, limit, source) return _pretty(data, pretty) @app.post("/search/parallel") async def search_parallel(req: BatchRequest): """Run up to 50 searches in parallel across DuckDB connection pool.""" # 🔑 KEY CHECK valid, err = _check_key(req.key) if not valid: return Response( content=json.dumps(err, ensure_ascii=False), status_code=403, media_type="application/json" ) if not req.queries: raise HTTPException(400, "queries must not be empty") if len(req.queries) > 50: raise HTTPException(400, "max 50 queries per batch") results = await _parallel([(item.get("field", "phoneNumber"), item.get("value", ""), item.get("mode", "exact"), int(item.get("limit", req.limit))) for item in req.queries]) out = [{"error": str(r)} if isinstance(r, Exception) else r for r in results] return _pretty({"searches": len(req.queries), "parallelism": PARALLELISM, "results": out}, True)