Update app.py

app.py

@@ -0,0 +1,299 @@
+"""FastAPI backend over DuckDB - ICMR + HITEK combined dataset."""
+import asyncio
+import json
+import os
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+from datetime import datetime
+
+import duckdb
+from fastapi import FastAPI, HTTPException, Query, Response
+from pydantic import BaseModel
+
+BASE = os.path.dirname(os.path.abspath(__file__))
+
+# ── Environment Variables (Secrets से आएँगी) ──────────
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+API_KEY = os.environ.get("API_KEY", "")  # 🔑 Secret से
+
+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)],
+}
+
+# 🔑 Valid Keys (secret से load)
+VALID_KEYS = {}
+if API_KEY:
+    VALID_KEYS[API_KEY] = {
+        "expires": 1830902400000,  # 6 November 2026
+    }
+
+app = FastAPI(title="ICMR + HITEK Search API",
+              description="DuckDB-backed search over 2.5B records")
+
+_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()
+    con.execute("INSTALL parquet; LOAD parquet;")
+    con.execute("INSTALL httpfs; LOAD httpfs;")
+
+    # 🔐 HF_TOKEN से dataset access (agar private hai)
+    if HF_TOKEN:
+        con.execute(f"SET http_headers = {{'Authorization': 'Bearer {HF_TOKEN}'}};")
+
+    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]
+
+
+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]]:
+    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("'", "''")
+
+    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": "@dmforbans"}
+    if not VALID_KEYS:
+        return False, {"message": "Provide correct key", "contact": "@dmforbans"}
+
+    key_info = VALID_KEYS.get(key)
+    if not key_info:
+        return False, {"message": "Provide correct key", "contact": "@dmforbans"}
+
+    now_ms = int(datetime.now().timestamp() * 1000)
+    if now_ms > key_info["expires"]:
+        expiry_date = datetime.fromtimestamp(key_info["expires"] / 1000).strftime("%Y-%m-%d")
+        return False, {
+            "message": "Your API key expired",
+            "expired_on": expiry_date,
+            "contact": "@dmforbans",
+        }
+
+    return True, key_info
+
+
+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")
+
+
+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]:
+    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": []}
+
+
+class BatchRequest(BaseModel):
+    queries: list[dict[str, Any]]
+    limit: int = 10
+    key: str | None = None
+
+
+# ---- ROUTES ----
+@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"),
+    field: str | None = Query(None, description=f"Restrict to one field"),
+    mode: str = Query("exact", pattern="^(exact|contains)$"),
+    limit: int = Query(10, ge=1, le=1000),
+    source: str | None = Query(None, pattern="^(icmr|hitek|inddata)$"),
+    pretty: bool = Query(True),
+):
+    # 🔑 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)
+
+    # ✅ Developer हमारा, credit हटाया
+    return _pretty({**data, "developer": "@dmforbans"}, pretty)
+
+
+@app.post("/search/parallel")
+async def search_parallel(req: BatchRequest):
+    # 🔑 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,
+                    "developer": "@dmforbans", "results": out}, True)