Create app.py

app.py

@@ -0,0 +1,304 @@
+"""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
+
+SEARCH_FIELDS = [
+    "name", "fathersName", "phoneNumber", "aadharNumber", "otherNumber",
+    "address", "district", "pincode", "state", "town", "source",
+]
+NUMBER_FIELDS = ["phoneNumber", "aadharNumber", "otherNumber"]
+
+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 = os.environ.get("API_KEY", "")
+
+VALID_KEYS = {}
+if API_KEY:
+    VALID_KEYS[API_KEY] = {}  # expiry नहीं — हमेशा चलेगी
+
+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}'}};")
+
+    # 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:
+        return False, {"message": "Provide correct key", "contact": "@asurhu"}
+    if not VALID_KEYS:
+        return False, {"message": "Provide correct key", "contact": "@dmforbans"}
+    if 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 @asurhu",
+    }
+
+
+@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)