Update app.py

app.py

@@ -1,5 +1,7 @@
 import os
-from fastapi import FastAPI, Header, HTTPException
+import duckdb
+
+from fastapi import FastAPI, Query, HTTPException
 
 app = FastAPI(
     title="InayatGaming Dataset API",
@@ -8,18 +10,33 @@ app = FastAPI(
 
 API_KEY = os.getenv("API_KEY")
 
+PARQUET_FILES = [
+    "hf://datasets/Inayatgaming/InayatGaming/part1.parquet",
+    "hf://datasets/Inayatgaming/InayatGaming/part2a.parquet",
+    "hf://datasets/Inayatgaming/InayatGaming/part2b_new.parquet",
+]
+
+# Only fields explicitly approved for API output.
+PUBLIC_FIELDS = [
+    "district",
+    "pincode",
+    "state",
+    "town",
+    "source",
+]
+
 
-def auth(key: str | None):
+def authenticate(api_key: str):
     if not API_KEY:
         raise HTTPException(
             status_code=500,
             detail="API_KEY is not configured"
         )
 
-    if key != API_KEY:
+    if api_key != API_KEY:
         raise HTTPException(
             status_code=401,
-            detail="Unauthorized"
+            detail="Invalid API key"
         )
 
 
@@ -33,22 +50,61 @@ def root():
 
 @app.get("/health")
 def health():
-    return {"status": "ok"}
+    return {
+        "status": "ok"
+    }
 
 
-@app.get("/info")
-def info(x_api_key: str | None = Header(default=None)):
-    auth(x_api_key)
+@app.get("/lookup")
+def lookup(
+    api_key: str = Query(...),
+    number: str = Query(..., min_length=1)
+):
+    authenticate(api_key)
 
-    return {
-        "dataset": "Inayatgaming/InayatGaming",
-        "format": "Parquet",
-        "files": [
-            "part1.parquet",
-            "part2a.parquet",
-            "part2b_new.parquet"
-        ],
-        "rows": 2504793870,
-        "approx_size": "103.8 GB",
-        "query_engine": "DuckDB"
-    }
+    # Normalize the supplied identifier.
+    number = number.strip()
+
+    con = duckdb.connect()
+
+    try:
+        query = """
+            SELECT
+                district,
+                pincode,
+                state,
+                town,
+                source
+            FROM read_parquet(
+                ?,
+                union_by_name=true
+            )
+            WHERE phoneNumber = ?
+            LIMIT 1
+        """
+
+        # Query each source file rather than loading the dataset
+        # into Python memory.
+        for parquet_file in PARQUET_FILES:
+            result = con.execute(
+                query,
+                [parquet_file, number]
+            ).fetchone()
+
+            if result:
+                return {
+                    "success": True,
+                    "match": True,
+                    "data": dict(
+                        zip(PUBLIC_FIELDS, result)
+                    )
+                }
+
+        return {
+            "success": True,
+            "match": False,
+            "data": None
+        }
+
+    finally:
+        con.close()