Update app.py

app.py

@@ -1,9 +1,13 @@
 import os
+import re
 import time
+import threading
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
 import streamlit as st
 import duckdb
 import pandas as pd
-from concurrent.futures import ThreadPoolExecutor, as_completed
 
 # ── Config ─────────────────────────────────────────────────────────────────────
 HF_TOKEN   = os.environ.get("HF_TOKEN", "")
@@ -14,6 +18,32 @@ _L2        = os.path.join(BASE_DIR, "mobile2_index.parquet")
 MOBILE1_IDX = _L1 if os.path.exists(_L1) else "hf://datasets/styhero/styhero_dataset/data/mobile1_index.parquet"
 MOBILE2_IDX = _L2 if os.path.exists(_L2) else "hf://datasets/styhero/styhero_dataset/data/mobile2_index.parquet"
 
+# Allowlist for column names that go into f-strings (can't be SQL params)
+_ALLOWED_COLS = {"mobile1", "mobile2"}
+
+# Mobile number validation: digits only, 10–15 chars (covers Indian + intl. format)
+_MOBILE_RE = re.compile(r"^\d{10,15}$")
+
+# ── Connection pool (per-thread reuse) ────────────────────────────────────────
+_local = threading.local()
+
+def get_con():
+    """Reuse one DuckDB connection per thread. Connection setup costs 50-150ms
+    so reusing it across queries in the same request is a big win."""
+    con = getattr(_local, "con", None)
+    if con is not None:
+        return con
+    con = duckdb.connect()
+    con.execute("INSTALL httpfs; LOAD httpfs;")
+    con.execute("CREATE SECRET IF NOT EXISTS hf_secret (TYPE HUGGINGFACE, TOKEN ?)", [HF_TOKEN])
+    con.execute("SET threads=4;")
+    con.execute("SET enable_progress_bar=false;")
+    con.execute("SET enable_object_cache=true;")   # cache parquet metadata across queries
+    con.execute("SET http_keep_alive=true;")
+    con.execute("SET http_retries=3;")
+    _local.con = con
+    return con
+
 # ── Helpers ────────────────────────────────────────────────────────────────────
 
 def fmt_ms(ms):
@@ -25,7 +55,6 @@ def clean(val):
     s = str(val).strip()
     if s.lower() in ("nan", "none", ""):
         return ""
-    # Strip trailing .0 from numeric values (DECIMAL stored as float)
     if s.endswith(".0"):
         try:
             s = str(int(float(s)))
@@ -40,50 +69,96 @@ def sanitize_df(df):
             df[col] = df[col].apply(clean)
     return df
 
-def new_con():
-    con = duckdb.connect()
-    con.execute("INSTALL httpfs; LOAD httpfs;")
-    con.execute("CREATE SECRET IF NOT EXISTS hf_secret (TYPE HUGGINGFACE, TOKEN ?)", [HF_TOKEN])
-    con.execute("SET threads=4; SET enable_progress_bar=false;")
-    return con
-
 # ── Queries ────────────────────────────────────────────────────────────────────
 
-def _query_one(index_path, col, value):
-    con = new_con()
-    t0  = time.perf_counter()
-    df  = con.execute(f"""
-        SELECT mobile1, mobile2, row_group_id, row_start, row_end, filename
-        FROM read_parquet('{index_path}')
-        WHERE {col} = '{value}'
-        ORDER BY mobile2
-    """).fetchdf()
+def query_both_indexes(target):
+    """Combined UNION ALL query over both index files in a single round trip.
+    Parameterized to prevent SQL injection."""
+    con = get_con()
+    t0 = time.perf_counter()
+    df = con.execute("""
+        SELECT mobile1, mobile2, row_group_id, row_start, row_end, filename, 'mobile1' AS src
+        FROM read_parquet(?)
+        WHERE mobile1 = ?
+        UNION ALL
+        SELECT mobile1, mobile2, row_group_id, row_start, row_end, filename, 'mobile2' AS src
+        FROM read_parquet(?)
+        WHERE mobile2 = ?
+    """, [MOBILE1_IDX, target, MOBILE2_IDX, target]).fetchdf()
     return df, (time.perf_counter() - t0) * 1000
 
-def query_both_parallel(target):
-    results = {}
-    with ThreadPoolExecutor(max_workers=2) as ex:
-        futures = {
-            ex.submit(_query_one, MOBILE1_IDX, "mobile1", target): "mobile1",
-            ex.submit(_query_one, MOBILE2_IDX, "mobile2", target): "mobile2",
-        }
-        for f in as_completed(futures):
-            df, ms = f.result()
-            results[futures[f]] = (df, ms)
-    return results
-
-def _fetch_one(filename, row_start, row_end, search_col, target):
-    con     = new_con()
+def _fetch_file_batch(filename, ranges, target):
+    """Fetch all matching row-groups for a single parquet file in one query.
+    `ranges` is a list of (row_start, row_end, search_col) tuples."""
+    # Validate search columns against allowlist before string interpolation
+    for _, _, col in ranges:
+        if col not in _ALLOWED_COLS:
+            raise ValueError(f"Invalid search column: {col!r}")
+
+    con = get_con()
     hf_path = f"{HF_DATASET}/{filename.split('/')[-1]}"
-    t0      = time.perf_counter()
-    df      = con.execute(f"""
+
+    # Build OR conditions; row_start/row_end are ints (we cast on input), col is allowlisted
+    or_parts = []
+    params = [hf_path]
+    for rs, re_, col in ranges:
+        or_parts.append(f"(file_row_number BETWEEN {int(rs)} AND {int(re_)} AND {col} = ?)")
+        params.append(target)
+
+    sql = f"""
         SELECT mobile1, mobile2, email, id, name, fname, address, circle
-        FROM read_parquet('{hf_path}', file_row_number=true)
-        WHERE file_row_number BETWEEN {row_start} AND {row_end}
-          AND {search_col} = '{target}'
-    """).fetchdf()
+        FROM read_parquet(?, file_row_number=true)
+        WHERE {' OR '.join(or_parts)}
+    """
+    t0 = time.perf_counter()
+    df = con.execute(sql, params).fetchdf()
     return sanitize_df(df), (time.perf_counter() - t0) * 1000
 
+# ── Cached search ─────────────────────────────────────────────────────────────
+
+@st.cache_data(ttl=300, max_entries=200, show_spinner=False)
+def search(target: str):
+    """Top-level cached search. Returns (combined_df, idx_ms, fetch_ms, errors).
+    TTL kept short because results contain PII — don't want them sitting in
+    memory indefinitely."""
+    errors = []
+    t_start = time.perf_counter()
+
+    idx_df, idx_ms = query_both_indexes(target)
+
+    if idx_df.empty:
+        return pd.DataFrame(), idx_ms, 0.0, errors
+
+    # Group row ranges by file so we make 1 HTTP query per file, not per row-group
+    by_file = defaultdict(list)
+    for _, rg in idx_df[["filename", "row_start", "row_end", "src"]].drop_duplicates().iterrows():
+        by_file[str(rg["filename"])].append(
+            (int(rg["row_start"]), int(rg["row_end"]), str(rg["src"]))
+        )
+
+    all_dfs = []
+    if by_file:
+        with ThreadPoolExecutor(max_workers=min(8, len(by_file))) as ex:
+            futures = {
+                ex.submit(_fetch_file_batch, fname, ranges, target): fname
+                for fname, ranges in by_file.items()
+            }
+            for f in as_completed(futures):
+                try:
+                    df, _ = f.result()
+                    if not df.empty:
+                        all_dfs.append(df)
+                except Exception as e:
+                    errors.append(str(e))
+
+    fetch_ms = (time.perf_counter() - t_start) * 1000
+
+    if not all_dfs:
+        return pd.DataFrame(), idx_ms, fetch_ms, errors
+
+    combined = pd.concat(all_dfs, ignore_index=True).drop_duplicates()
+    return combined, idx_ms, fetch_ms, errors
+
 # ── SVG icons ──────────────────────────────────────────────────────────────────
 
 def ico(path, color="#8b8bff", size=14):
@@ -96,58 +171,6 @@ ICO_PIN     = ico('<path d="M20 10c0 6-8 12-8 12s-8-6-8-12a8 8 0 0 1 16 0Z"/><ci
 ICO_CIRCLE  = ico('<circle cx="12" cy="12" r="10"/><path d="M12 8v4l3 3"/>')
 ICO_PERSON  = ico('<path d="M20 21v-2a4 4 0 0 0-4-4H8a4 4 0 0 0-4 4v2"/><circle cx="12" cy="7" r="4"/>')
 
-# ── Skeleton HTML ──────────────────────────────────────────────────────────────
-
-SKELETON_CSS = """
-<style>
-@keyframes shimmer {
-  0%   { background-position: -600px 0 }
-  100% { background-position:  600px 0 }
-}
-.sk {
-  border-radius: 6px;
-  background: linear-gradient(90deg,rgba(255,255,255,.05) 25%,rgba(255,255,255,.12) 50%,rgba(255,255,255,.05) 75%);
-  background-size: 1200px 100%;
-  animation: shimmer 1.5s infinite linear;
-}
-.sk-card {
-  background: rgba(255,255,255,.05);
-  border: 1px solid rgba(255,255,255,.1);
-  border-radius: 16px;
-  padding: 20px;
-  margin-bottom: 14px;
-}
-.sk-row { display:flex; gap:12px; align-items:center; margin-bottom:10px; }
-.sk-circle { width:36px;height:36px;border-radius:50%; flex-shrink:0; }
-.sk-line { height:14px; }
-.sk-grid { display:grid; grid-template-columns:1fr 1fr; gap:12px; margin-top:14px; }
-.sk-field { display:flex; flex-direction:column; gap:5px; }
-.sk-label { height:9px; width:40%; }
-.sk-value { height:13px; width:75%; }
-</style>
-"""
-
-def skeleton_card():
-    return """
-<div class="sk-card">
-  <div class="sk-row">
-    <div class="sk sk-circle"></div>
-    <div style="flex:1;display:flex;flex-direction:column;gap:6px">
-      <div class="sk sk-line" style="width:55%"></div>
-      <div class="sk sk-line" style="width:35%;height:10px"></div>
-    </div>
-  </div>
-  <div class="sk-grid">
-    <div class="sk-field"><div class="sk sk-label"></div><div class="sk sk-value"></div></div>
-    <div class="sk-field"><div class="sk sk-label"></div><div class="sk sk-value"></div></div>
-    <div class="sk-field"><div class="sk sk-label"></div><div class="sk sk-value" style="width:85%"></div></div>
-    <div class="sk-field"><div class="sk sk-label"></div><div class="sk sk-value" style="width:60%"></div></div>
-    <div class="sk-field" style="grid-column:1/-1"><div class="sk sk-label"></div><div class="sk sk-value" style="width:90%"></div></div>
-  </div>
-</div>"""
-
-SKELETONS_HTML = SKELETON_CSS + "".join(skeleton_card() for _ in range(3))
-
 # ── Result card ────────────────────────────────────────────────────────────────
 
 def field_html(ico_svg, label, value, full=False):
@@ -269,7 +292,6 @@ const resize = () => {{
     if (f.contentWindow === window) f.style.height = (h + 8) + 'px';
   }});
 }};
-// ResizeObserver fires on layout shifts, font load, grid reflow
 if (window.ResizeObserver) {{
   new ResizeObserver(resize).observe(card);
 }}
@@ -288,7 +310,6 @@ st.set_page_config(page_title="Mobile Lookup", layout="centered")
 
 st.markdown("""
 <style>
-/* ── Reset & background ─────────────────────────────────────────────────── */
 .stApp, .stApp > div {
   background: linear-gradient(145deg,#060612 0%,#0e0b24 40%,#16103a 100%) !important;
   background-attachment: fixed !important;
@@ -300,7 +321,6 @@ st.markdown("""
   max-width: 680px !important;
 }
 
-/* ── Remove ALL Streamlit input wrappers ────────────────────────────────── */
 .stTextInput > div,
 .stTextInput > div > div {
   background: transparent !important;
@@ -309,7 +329,6 @@ st.markdown("""
   padding: 0 !important;
 }
 
-/* Single clean input */
 .stTextInput input {
   background: rgba(255,255,255,.06) !important;
   border: 1px solid rgba(255,255,255,.12) !important;
@@ -330,7 +349,6 @@ st.markdown("""
 .stTextInput label { display: none !important; }
 [data-testid="InputInstructions"] { display: none !important; }
 
-/* ── Form glass panel ───────────────────────────────────────────────────── */
 [data-testid="stForm"] {
   background: rgba(255,255,255,.04) !important;
   backdrop-filter: blur(24px) !important;
@@ -341,7 +359,6 @@ st.markdown("""
   box-shadow: 0 20px 60px rgba(0,0,0,.5), inset 0 1px 0 rgba(255,255,255,.07) !important;
 }
 
-/* ── Submit button ──────────────────────────────────────────────────────── */
 .stFormSubmitButton > button {
   background: linear-gradient(135deg,#6c5ce7 0%,#a855f7 100%) !important;
   color: #fff !important;
@@ -362,13 +379,11 @@ st.markdown("""
 }
 .stFormSubmitButton > button:active { transform: translateY(0) !important; }
 
-/* ── Spinner ────────────────────────────────────────────────────────────── */
 .stSpinner > div {
   border-color: rgba(255,255,255,.1) !important;
   border-top-color: #7c6fff !important;
 }
 
-/* ── Alerts ─────────────────────────────────────────────────────────────── */
 [data-testid="stAlert"] {
   background: rgba(255,255,255,.04) !important;
   border: 1px solid rgba(255,255,255,.08) !important;
@@ -377,12 +392,10 @@ st.markdown("""
   color: #d0d0f0 !important;
 }
 
-/* ── Scrollbar ──────────────────────────────────────────────────────────── */
 ::-webkit-scrollbar { width: 6px; }
 ::-webkit-scrollbar-track { background: transparent; }
 ::-webkit-scrollbar-thumb { background: rgba(255,255,255,.1); border-radius: 3px; }
 
-/* ── iframe card spacing ─────────────────────────────────────────────────── */
 iframe { display: block !important; }
 [data-testid="stCustomComponentV1"] { margin-bottom: 0 !important; line-height: 0; }
 </style>
@@ -416,8 +429,11 @@ st.markdown("""
 # ── Search form ─────────────────────────────────────────────────────────────────
 
 with st.form("search_form"):
-    number = st.text_input("", placeholder="Enter mobile number — e.g. 919876543210",
-                           label_visibility="collapsed")
+    number = st.text_input(
+        "Mobile number",
+        placeholder="Enter mobile number — e.g. 919876543210",
+        label_visibility="collapsed",
+    )
     submitted = st.form_submit_button("Search", use_container_width=True, type="primary")
 
 if not submitted or not number.strip():
@@ -425,62 +441,28 @@ if not submitted or not number.strip():
 
 target = number.strip()
 
-# ── Search with progress ────────────────────────────────────────────────────────
+# Input validation — reject anything that isn't 10-15 digits
+if not _MOBILE_RE.match(target):
+    st.error("Please enter a valid mobile number (10–15 digits, numbers only).")
+    st.stop()
 
-all_dfs, errors = [], []
-t0 = time.perf_counter()
+# ── Search with progress ────────────────────────────────────────────────────────
 
 with st.status("Searching...", expanded=False) as status:
+    combined, idx_ms, fetch_ms, errors = search(target)
 
-    idx_results  = query_both_parallel(target)
-    idx_ms       = (time.perf_counter() - t0) * 1000
-    m1_df, m1_ms = idx_results["mobile1"]
-    m2_df, m2_ms = idx_results["mobile2"]
-
-    if m1_df.empty and m2_df.empty:
+    if combined.empty:
         status.update(label="No records found", state="error", expanded=False)
+        for err in errors:
+            st.error(err)
         st.warning(f"No records found for **{target}**")
         st.stop()
 
-    fetch_tasks = []
-    for src_col, src_df in [("mobile1", m1_df), ("mobile2", m2_df)]:
-        if src_df.empty:
-            continue
-        for _, rg in src_df[["filename","row_group_id","row_start","row_end"]].drop_duplicates().iterrows():
-            fetch_tasks.append({
-                "filename":   str(rg["filename"]),
-                "row_start":  int(rg["row_start"]),
-                "row_end":    int(rg["row_end"]),
-                "search_col": src_col,
-            })
-
-    status.update(label="Loading records...", state="running", expanded=False)
-
-    with ThreadPoolExecutor(max_workers=min(8, len(fetch_tasks))) as ex:
-        futures = {
-            ex.submit(_fetch_one, t["filename"], t["row_start"], t["row_end"], t["search_col"], target): t
-            for t in fetch_tasks
-        }
-        for f in as_completed(futures):
-            try:
-                df, _ = f.result()
-                if not df.empty:
-                    all_dfs.append(df)
-            except Exception as e:
-                errors.append(str(e))
-
-    fetch_ms = (time.perf_counter() - t0) * 1000
     status.update(label=f"Done — {fmt_ms(fetch_ms)}", state="complete", expanded=False)
 
 for err in errors:
     st.error(err)
 
-if not all_dfs:
-    st.warning("No matching rows found in remote parquet.")
-    st.stop()
-
-combined = pd.concat(all_dfs, ignore_index=True).drop_duplicates()
-
 # ── Results ─────────────────────────────────────────────────────────────────────
 
 st.markdown(f"""
@@ -497,4 +479,4 @@ st.markdown(f"""
 """, unsafe_allow_html=True)
 
 for i, (_, row) in enumerate(combined.iterrows()):
-    show_card(row, i)
+    show_card(row, i)