screen-leads-app/db/database.py

"""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()