"""SQLite persistence layer: companies table, jobs table, dedup, company cache, CSV export.""" from __future__ import annotations import csv import logging import sqlite3 from datetime import datetime, timezone from pathlib import Path from urllib.parse import urlsplit from .config import get_settings from .models import CSV_COLUMNS, JobResult, JobStatus logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Pure helper — standalone so tests can import it directly # --------------------------------------------------------------------------- def _company_key(name: str, website: str | None) -> str: """Derive a stable company key: registrable domain if website known, else name. Examples: _company_key("Acme Inc", "https://www.acme.com/about") -> "acme.com" _company_key("Acme LLC", None) -> "acme" """ if website: try: hostname = urlsplit(website).hostname or "" hostname = hostname.lower() if hostname.startswith("www."): hostname = hostname[4:] if hostname: return hostname except Exception: pass slug = name.strip().lower() return slug or "unknown" def _now_iso() -> str: return datetime.now(timezone.utc).isoformat() # --------------------------------------------------------------------------- # Database # --------------------------------------------------------------------------- class Database: """SQLite-backed store for companies and jobs. Supports use as a context manager. Thread-safety: one connection per instance (single-threaded pipeline use). Postgres-swap is possible by replacing this class behind the same interface. """ def __init__(self, path: Path | str | None = None) -> None: if path is None: path = get_settings().db_path path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) self._conn = sqlite3.connect(str(path)) self._conn.row_factory = sqlite3.Row # WAL mode for better write performance (no-op on in-memory dbs) self._conn.execute("PRAGMA journal_mode=WAL") self.init_schema() logger.debug("db: opened %s", path) # ---- Schema ----------------------------------------------------------- def init_schema(self) -> None: """Create tables if they don't exist. Safe to call on every startup.""" self._conn.executescript(""" CREATE TABLE IF NOT EXISTS companies ( company_key TEXT PRIMARY KEY, name TEXT, website TEXT, career_url TEXT, first_seen TEXT ); CREATE TABLE IF NOT EXISTS jobs ( job_id TEXT PRIMARY KEY, company_key TEXT, linkedin_url TEXT, position_url TEXT, status TEXT, listed_at TEXT, first_seen TEXT ); """) self._conn.commit() # ---- Dedup / cache ---------------------------------------------------- def is_seen(self, job_id: str) -> bool: """Return True if this job_id already exists in the DB (cross-run dedup).""" row = self._conn.execute( "SELECT 1 FROM jobs WHERE job_id = ?", (job_id,) ).fetchone() return row is not None def get_cached_career_url(self, company_key: str) -> str | None: """Return a previously resolved career URL for this company, or None. When non-null, the pipeline skips the cascade for this company and uses the cached value directly (companies are resolved at most once). """ row = self._conn.execute( "SELECT career_url FROM companies WHERE company_key = ?", (company_key,) ).fetchone() if row and row["career_url"]: return row["career_url"] return None # ---- Company key helper ----------------------------------------------- def company_key_for(self, name: str, website: str | None) -> str: """Delegate to the module-level pure function.""" return _company_key(name, website) # ---- Upserts ---------------------------------------------------------- def upsert_company( self, company_key: str, name: str, website: str | None, career_url: str | None, ) -> None: """Insert or update a company row. COALESCE semantics: a new non-null value fills an existing NULL, but never clobbers an already-populated column. first_seen is preserved. """ now = _now_iso() self._conn.execute( """ INSERT INTO companies(company_key, name, website, career_url, first_seen) VALUES (?, ?, ?, ?, ?) ON CONFLICT(company_key) DO UPDATE SET name = COALESCE(excluded.name, companies.name), website = COALESCE(excluded.website, companies.website), career_url = COALESCE(excluded.career_url, companies.career_url), first_seen = companies.first_seen """, (company_key, name, website, career_url, now), ) self._conn.commit() def upsert_job(self, result: JobResult) -> None: """Insert or update a job row. first_seen is preserved on conflict. position_url and linkedin_url use COALESCE: a resolved value fills a prior NULL but is never overwritten once set. """ now = _now_iso() listed_at = ( result.listed_at.isoformat() if result.listed_at is not None else None ) self._conn.execute( """ INSERT INTO jobs(job_id, company_key, linkedin_url, position_url, status, listed_at, first_seen) VALUES (?, ?, ?, ?, ?, ?, ?) ON CONFLICT(job_id) DO UPDATE SET company_key = excluded.company_key, linkedin_url = COALESCE(excluded.linkedin_url, jobs.linkedin_url), position_url = COALESCE(excluded.position_url, jobs.position_url), status = excluded.status, listed_at = COALESCE(excluded.listed_at, jobs.listed_at), first_seen = jobs.first_seen """, ( result.job_id, result.company_key, result.linkedin_url, result.open_position_url, result.status.value, listed_at, now, ), ) self._conn.commit() def persist_result(self, result: JobResult) -> None: """Persist one JobResult atomically: upsert company then job.""" with self._conn: now = _now_iso() listed_at = ( result.listed_at.isoformat() if result.listed_at is not None else None ) self._conn.execute( """ INSERT INTO companies(company_key, name, website, career_url, first_seen) VALUES (?, ?, ?, ?, ?) ON CONFLICT(company_key) DO UPDATE SET name = COALESCE(excluded.name, companies.name), website = COALESCE(excluded.website, companies.website), career_url = COALESCE(excluded.career_url, companies.career_url), first_seen = companies.first_seen """, ( result.company_key, result.company_name, result.website, result.career_page_url, now, ), ) self._conn.execute( """ INSERT INTO jobs(job_id, company_key, linkedin_url, position_url, status, listed_at, first_seen) VALUES (?, ?, ?, ?, ?, ?, ?) ON CONFLICT(job_id) DO UPDATE SET company_key = excluded.company_key, linkedin_url = COALESCE(excluded.linkedin_url, jobs.linkedin_url), position_url = COALESCE(excluded.position_url, jobs.position_url), status = excluded.status, listed_at = COALESCE(excluded.listed_at, jobs.listed_at), first_seen = jobs.first_seen """, ( result.job_id, result.company_key, result.linkedin_url, result.open_position_url, result.status.value, listed_at, now, ), ) # ---- Export ----------------------------------------------------------- def export_csv(self, path: Path | None = None) -> Path: """Write output/results.csv; complete (position_found) rows sorted first. Uses models.CSV_COLUMNS for the header; None values become empty strings. Returns the Path written. """ if path is None: path = get_settings().output_csv path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) rows = self._conn.execute( """ SELECT c.name AS company_name, c.career_url AS career_page_url, j.position_url AS open_position_url, j.status FROM jobs j LEFT JOIN companies c ON j.company_key = c.company_key ORDER BY (j.status = 'position_found') DESC, c.name ASC NULLS LAST """ ).fetchall() with path.open("w", newline="", encoding="utf-8") as fh: writer = csv.DictWriter(fh, fieldnames=list(CSV_COLUMNS)) writer.writeheader() for row in rows: writer.writerow( { "company_name": row["company_name"] or "", "career_page_url": row["career_page_url"] or "", "open_position_url": row["open_position_url"] or "", } ) logger.info("db: exported %d rows to %s", len(rows), path) return path # ---- Counts (for summary / tests) ------------------------------------ def counts(self) -> dict[str, int]: """Return per-status job counts from the DB.""" rows = self._conn.execute( "SELECT status, COUNT(*) AS n FROM jobs GROUP BY status" ).fetchall() return {row["status"]: row["n"] for row in rows} # ---- Lifecycle ------------------------------------------------------- def close(self) -> None: self._conn.close() def __enter__(self) -> "Database": return self def __exit__(self, *_: object) -> None: self.close()