"""SQLite persistence for leads and capture runs. Thread-safe: the capture loop runs on a background thread while the API serves on the main thread, so all access goes through one connection guarded by a lock. Leads are keyed by source_url; re-capturing the same profile merges new, non-empty fields into the existing row rather than duplicating it. """ from __future__ import annotations import json import sqlite3 import threading import time from typing import Any, Optional from config import DB_PATH from core.models import Lead _lock = threading.Lock() _conn: Optional[sqlite3.Connection] = None _JSON_FIELDS = ("social_handles", "screenshot_refs", "fields_found") def _connect() -> sqlite3.Connection: global _conn if _conn is None: _conn = sqlite3.connect(str(DB_PATH), check_same_thread=False) _conn.row_factory = sqlite3.Row return _conn def init_db() -> None: with _lock: conn = _connect() conn.executescript( """ CREATE TABLE IF NOT EXISTS leads ( id INTEGER PRIMARY KEY AUTOINCREMENT, full_name TEXT, first_name TEXT, last_name TEXT, headline TEXT, title TEXT, company TEXT, company_url TEXT, industry TEXT, location TEXT, about TEXT, email TEXT, phone TEXT, website TEXT, social_handles TEXT, source_site TEXT, source_url TEXT UNIQUE, captured_at REAL, updated_at REAL, screenshot_refs TEXT, confidence REAL, fields_found TEXT, needs_review INTEGER, stage TEXT ); CREATE TABLE IF NOT EXISTS runs ( id INTEGER PRIMARY KEY AUTOINCREMENT, status TEXT, started_at REAL, stopped_at REAL, reason TEXT, leads_captured INTEGER DEFAULT 0 ); """ ) conn.commit() def _row_to_lead(row: sqlite3.Row) -> dict[str, Any]: d = dict(row) for f in _JSON_FIELDS: try: d[f] = json.loads(d[f]) if d.get(f) else ([] if f != "social_handles" else {}) except Exception: d[f] = [] if f != "social_handles" else {} d["needs_review"] = bool(d.get("needs_review")) return d def get_lead_by_url(source_url: str) -> Optional[dict[str, Any]]: with _lock: cur = _connect().execute("SELECT * FROM leads WHERE source_url = ?", (source_url,)) row = cur.fetchone() return _row_to_lead(row) if row else None def upsert_lead(lead: Lead) -> dict[str, Any]: """Insert a new lead, or merge non-empty fields into an existing one.""" lead.finalize() with _lock: conn = _connect() existing = conn.execute( "SELECT * FROM leads WHERE source_url = ?", (lead.source_url,) ).fetchone() d = lead.to_dict() for f in _JSON_FIELDS: d[f] = json.dumps(d[f], ensure_ascii=False) d["needs_review"] = 1 if d["needs_review"] else 0 if existing is None: cols = ", ".join(d.keys()) ph = ", ".join(["?"] * len(d)) conn.execute(f"INSERT INTO leads ({cols}) VALUES ({ph})", tuple(d.values())) else: merged = dict(existing) # prefer new non-empty scalar values; JSON fields already serialised for k, v in d.items(): if k in ("id", "captured_at", "source_url"): continue if k in _JSON_FIELDS: if v not in ("[]", "{}", "", None): merged[k] = v elif v not in ("", None, 0, 0.0): merged[k] = v merged["updated_at"] = time.time() set_clause = ", ".join(f"{k} = ?" for k in merged if k != "id") vals = [merged[k] for k in merged if k != "id"] conn.execute( f"UPDATE leads SET {set_clause} WHERE id = ?", (*vals, merged["id"]) ) conn.commit() return get_lead_by_url(lead.source_url) or {} def list_leads(stage: Optional[str] = None) -> list[dict[str, Any]]: with _lock: conn = _connect() if stage: cur = conn.execute( "SELECT * FROM leads WHERE stage = ? ORDER BY updated_at DESC", (stage,) ) else: cur = conn.execute("SELECT * FROM leads ORDER BY updated_at DESC") return [_row_to_lead(r) for r in cur.fetchall()] def stage_counts() -> dict[str, int]: with _lock: cur = _connect().execute("SELECT stage, COUNT(*) c FROM leads GROUP BY stage") return {r["stage"]: r["c"] for r in cur.fetchall()} def set_stage(lead_id: int, stage: str) -> None: with _lock: conn = _connect() conn.execute( "UPDATE leads SET stage = ?, updated_at = ? WHERE id = ?", (stage, time.time(), lead_id), ) conn.commit() # --- runs -------------------------------------------------------------------- def start_run() -> int: with _lock: conn = _connect() cur = conn.execute( "INSERT INTO runs (status, started_at) VALUES ('running', ?)", (time.time(),) ) conn.commit() return int(cur.lastrowid) def finish_run(run_id: int, status: str, reason: str, leads_captured: int) -> None: with _lock: conn = _connect() conn.execute( "UPDATE runs SET status = ?, stopped_at = ?, reason = ?, leads_captured = ? WHERE id = ?", (status, time.time(), reason, leads_captured, run_id), ) conn.commit()