phase5-pipeline
This commit is contained in:
305
jobsource/db.py
305
jobsource/db.py
@@ -1,10 +1,297 @@
|
||||
"""SQLite persistence layer: companies table, jobs table, dedup, company cache, CSV export.
|
||||
"""SQLite persistence layer: companies table, jobs table, dedup, company cache, CSV export."""
|
||||
from __future__ import annotations
|
||||
|
||||
Scaffold stub -- not implemented yet.
|
||||
"""
|
||||
# TODO (Stage 4): implement per CLAUDE.md "Stage 4 — Persist & export" and "Data model".
|
||||
# Schema:
|
||||
# companies(company_key PK, name, website, career_url, first_seen)
|
||||
# jobs(job_id PK, company_key, linkedin_url, position_url, status, listed_at, first_seen)
|
||||
# CSV export writes output/results.csv with columns: company_name, career_page_url, open_position_url
|
||||
# (complete rows — status==position_found — sorted first; incomplete rows follow).
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user