298 lines
11 KiB
Python
298 lines
11 KiB
Python
"""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()
|