diff --git a/app/services/data_import.py b/app/services/data_import.py new file mode 100644 index 0000000..494ffb3 --- /dev/null +++ b/app/services/data_import.py @@ -0,0 +1,255 @@ +"""Source-agnostic batch import framework (Dolt/SEC bulk data → PostgreSQL). + +Every bulk importer (SEC facts, Dolt earnings, later Dolt stocks) plugs into +``run_import`` and gets, for free, the plan's non-negotiables: + +- **One run per source at a time** — a Postgres *session-level* advisory lock + keyed by source. It is held on a single pinned connection for the whole run, + so it survives the intermediate commits (the ``running`` row, then the + promotion) and only releases at the end. No-op on non-Postgres (tests). +- **Idempotent per revision** — the cheap ``detect_revision`` probe is compared + against the last *promoted* run; an unchanged revision records a ``no_op`` + with **zero row changes** (no expensive fetch, no writes). +- **Staging then atomic promotion** — the importer stages into an in-memory + object (no physical staging tables), validation reads it, and only a passing + run calls ``promote`` whose writes commit together with the run-row flip to + ``promoted`` in a single transaction. +- **Failure is inert** — a failed validation or a mid-run exception marks the + run ``failed``, alerts via the system-events path, and leaves the live tables + exactly as they were (nothing is written before ``promote``). + +Every attempt — promoted, no_op, or failed — is recorded in ``data_import_runs``. +KISS: no conflicts table (summaries go in ``validation_json``), no revision +table (idempotency queries the last run), no aggregate tables. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +from dataclasses import dataclass, field +from datetime import date, datetime, timezone +from typing import Any, Protocol, runtime_checkable + +from sqlalchemy import select, text +from sqlalchemy.engine import Engine # noqa: F401 (typing only) +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession + +from app.database import engine as app_engine +from app.models.data_import_run import DataImportRun +from app.services import system_event_service + +logger = logging.getLogger(__name__) + +# data_import_runs.status values +STATUS_RUNNING = "running" +STATUS_VALIDATED = "validated" +STATUS_PROMOTED = "promoted" +STATUS_NO_OP = "no_op" +STATUS_FAILED = "failed" + +_MAX_ERROR_LEN = 4000 + + +@dataclass +class ValidationResult: + """Outcome of an importer's validation gates. + + ``summary`` is serialized into ``validation_json`` (reconciliation / + discrepancy details live here — no separate conflicts table). ``validate`` + MUST be read-only: it reads the staged object and, if needed, live tables + for comparison, but writes nothing — that invariant is what makes a failed + run leave the dataset untouched. + """ + + ok: bool + summary: dict[str, Any] = field(default_factory=dict) + source_max_date: date | None = None + messages: list[str] = field(default_factory=list) + + +@runtime_checkable +class SourceImporter(Protocol): + """Interface a concrete bulk importer implements. All methods receive the + session bound to the lock-holding connection; ``stage`` and ``validate`` + never write to live tables, only ``promote`` does.""" + + source: str # sec_facts | dolt_earnings | dolt_stocks + + async def detect_revision(self, db: AsyncSession) -> str | None: + """Cheap probe of the source revision (Dolt commit / SEC archive SHA). + + Returns the revision id, or None when it can't be determined cheaply + (in which case idempotency is skipped and the run always stages).""" + ... + + async def stage(self, db: AsyncSession) -> Any: + """Download/parse into an in-memory staged representation. No writes to + live tables.""" + ... + + async def validate(self, db: AsyncSession, staged: Any) -> ValidationResult: + """Run the source's validation gates against ``staged``. Read-only.""" + ... + + async def promote(self, db: AsyncSession, staged: Any) -> dict[str, int]: + """Apply ``staged`` to the live tables. Called inside the promotion + transaction; the caller commits. Returns row-count deltas.""" + ... + + +def _advisory_key(source: str) -> int: + """Deterministic signed 64-bit key for a source's advisory lock.""" + digest = hashlib.blake2b(source.encode("utf-8"), digest_size=8).digest() + return int.from_bytes(digest, "big", signed=True) + + +async def _last_promoted_revision(db: AsyncSession, source: str) -> str | None: + """Revision of the most recent *promoted* run for ``source`` (the revision + currently loaded), or None if none has promoted yet.""" + row = await db.execute( + select(DataImportRun.revision) + .where( + DataImportRun.source == source, + DataImportRun.status == STATUS_PROMOTED, + ) + .order_by(DataImportRun.id.desc()) + .limit(1) + ) + return row.scalar_one_or_none() + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +async def _alert(db: AsyncSession, source: str, code: str, messages: list[str]) -> None: + try: + await system_event_service.log_event( + db, + severity="error", + source="data_import", + code=f"{source}_{code}", + message=(("; ".join(messages)) or code)[:_MAX_ERROR_LEN], + dedup_key=f"data_import:{source}:{code}", + ) + except Exception: # noqa: BLE001 — alerting must never mask the real outcome + logger.exception("Failed to emit data_import alert %s/%s", source, code) + + +async def run_import( + importer: SourceImporter, + *, + engine: AsyncEngine | None = None, +) -> DataImportRun | None: + """Run one import for ``importer``. + + Returns the recorded ``DataImportRun`` (promoted / no_op / failed), or None + when the per-source advisory lock is already held (another run is active). + """ + engine = engine or app_engine + source = importer.source + is_pg = engine.dialect.name == "postgresql" + key = _advisory_key(source) + + async with engine.connect() as conn: + # Bind the session to this one connection so the session-level advisory + # lock persists across our commits. expire_on_commit must be set here — + # the app factory's setting doesn't carry to a directly-built session. + session = AsyncSession(bind=conn, expire_on_commit=False) + try: + if is_pg: + got = ( + await session.execute( + text("SELECT pg_try_advisory_lock(:k)"), {"k": key} + ) + ).scalar() + await session.commit() + if not got: + logger.info("data_import %s: lock held, skipping", source) + return None + + revision = await importer.detect_revision(session) + last_rev = await _last_promoted_revision(session, source) + if revision is not None and revision == last_rev: + run = DataImportRun( + source=source, + revision=revision, + status=STATUS_NO_OP, + started_at=_now(), + completed_at=_now(), + ) + session.add(run) + await session.commit() + await session.refresh(run) + logger.info("data_import %s: no_op (revision %s)", source, revision) + return run + + run = DataImportRun( + source=source, + revision=revision, + status=STATUS_RUNNING, + started_at=_now(), + ) + session.add(run) + await session.commit() + await session.refresh(run) + + try: + staged = await importer.stage(session) + result = await importer.validate(session, staged) + run.source_max_date = result.source_max_date + run.validation_json = json.dumps(result.summary, default=str) + + if not result.ok: + run.status = STATUS_FAILED + run.error_details = ("; ".join(result.messages))[:_MAX_ERROR_LEN] + run.completed_at = _now() + await session.commit() + await _alert(session, source, "validation_failed", result.messages) + logger.warning( + "data_import %s: validation failed: %s", + source, + result.messages, + ) + return run + + # Promotion: importer writes + run-row flip in one transaction. + row_counts = await importer.promote(session, staged) + run.status = STATUS_PROMOTED + run.row_counts_json = json.dumps(row_counts, default=str) + run.completed_at = _now() + await session.commit() + await session.refresh(run) + logger.info( + "data_import %s: promoted (revision %s, rows %s)", + source, + revision, + row_counts, + ) + return run + + except Exception as exc: # noqa: BLE001 — record + alert, don't crash the job + await session.rollback() + run.status = STATUS_FAILED + run.error_details = repr(exc)[:_MAX_ERROR_LEN] + run.completed_at = _now() + try: + await session.commit() + except Exception: # noqa: BLE001 + logger.exception("data_import %s: failed to record failure", source) + await _alert(session, source, "import_error", [repr(exc)]) + logger.exception("data_import %s: import error", source) + return run + + finally: + if is_pg: + try: + await session.execute( + text("SELECT pg_advisory_unlock(:k)"), {"k": key} + ) + await session.commit() + except Exception: # noqa: BLE001 + logger.exception("data_import %s: failed to release lock", source) + await session.close() diff --git a/tests/unit/test_data_import_framework.py b/tests/unit/test_data_import_framework.py new file mode 100644 index 0000000..9a30141 --- /dev/null +++ b/tests/unit/test_data_import_framework.py @@ -0,0 +1,201 @@ +"""Orchestration tests for the source-agnostic import framework. + +These drive ``run_import`` with a fake importer to prove the framework's +guarantees: idempotent no_op on an unchanged revision, atomic promotion on a +new revision, and — the load-bearing one — a failed validation or a mid-run +exception leaves the live tables untouched. + +The advisory-lock branch is a no-op on SQLite, so lock mutual-exclusion has NO +coverage here (PG-verify-pending); only the deterministic key derivation is +unit-tested. Per the SQLite StaticPool caveat we never share a connection: each +test uses its own temp-file engine and seeds/asserts with short-lived sessions +sequenced around the ``run_import`` call. +""" + +from __future__ import annotations + +import os +import tempfile +from datetime import date, datetime, timezone + +import pytest +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from app.database import Base +import app.models # noqa: F401 register models on Base.metadata +from app.models.data_import_run import DataImportRun +from app.models.fundamental_snapshot import FundamentalSnapshot +from app.models.system_event import SystemEvent +from app.services.data_import import ( + STATUS_FAILED, + STATUS_NO_OP, + STATUS_PROMOTED, + ValidationResult, + _advisory_key, + run_import, +) + + +@pytest.fixture +async def engine(): + """A dedicated temp-file SQLite engine (independent connections, unlike the + shared in-memory test engine) so ``run_import`` can pin its own connection.""" + fd, path = tempfile.mkstemp(suffix=".db") + os.close(fd) + eng = create_async_engine(f"sqlite+aiosqlite:///{path}") + async with eng.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + try: + yield eng + finally: + await eng.dispose() + try: + os.unlink(path) + except OSError: + pass + + +def _factory(engine) -> async_sessionmaker[AsyncSession]: + return async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) + + +class FakeImporter: + """Writes ``n_rows`` fundamental_snapshots on promote (CIK-keyed, no ticker + FK) so the live-table effect is easy to count.""" + + source = "sec_facts" + + def __init__(self, revision, *, ok=True, n_rows=3, raise_in="none"): + self.revision = revision + self.ok = ok + self.n_rows = n_rows + self.raise_in = raise_in + self.staged_called = False + self.promoted = False + + async def detect_revision(self, db): + if self.raise_in == "detect": + raise RuntimeError("boom-detect") + return self.revision + + async def stage(self, db): + self.staged_called = True + if self.raise_in == "stage": + raise RuntimeError("boom-stage") + return list(range(self.n_rows)) + + async def validate(self, db, staged): + return ValidationResult( + ok=self.ok, + summary={"staged_rows": len(staged)}, + source_max_date=date(2026, 7, 21), + messages=[] if self.ok else ["coverage below threshold"], + ) + + async def promote(self, db, staged): + if self.raise_in == "promote": + # write one row THEN raise, to prove rollback undoes partial writes + db.add(_snapshot(self.revision, 999)) + raise RuntimeError("boom-promote") + for i in staged: + db.add(_snapshot(self.revision, i)) + self.promoted = True + return {"fundamental_snapshots": len(staged)} + + +def _snapshot(revision: str, i: int) -> FundamentalSnapshot: + return FundamentalSnapshot( + cik=f"{i:010d}", + accession=f"{revision}-{i:06d}", + form="10-Q", + filed_date=date(2026, 7, 1), + accepted_at=datetime(2026, 7, 1, tzinfo=timezone.utc), + period_end=date(2026, 6, 30), + fiscal_year=2026, + fiscal_period="Q2", + revenue=1000.0 + i, + ) + + +async def _count(factory, model) -> int: + async with factory() as s: + return (await s.execute(select(func.count()).select_from(model))).scalar_one() + + +async def _runs(factory) -> list[DataImportRun]: + async with factory() as s: + return list( + (await s.execute(select(DataImportRun).order_by(DataImportRun.id))).scalars() + ) + + +# --------------------------------------------------------------------------- + + +def test_advisory_key_deterministic_and_distinct(): + assert _advisory_key("sec_facts") == _advisory_key("sec_facts") + assert _advisory_key("sec_facts") != _advisory_key("dolt_earnings") + for src in ("sec_facts", "dolt_earnings", "dolt_stocks"): + k = _advisory_key(src) + assert -(2**63) <= k < 2**63 # fits Postgres bigint + + +async def test_promote_writes_and_records_run(engine): + factory = _factory(engine) + run = await run_import(FakeImporter("rev1", n_rows=4), engine=engine) + + assert run is not None and run.status == STATUS_PROMOTED + assert run.revision == "rev1" + assert run.row_counts_json is not None and "fundamental_snapshots" in run.row_counts_json + assert run.source_max_date == date(2026, 7, 21) + assert run.completed_at is not None + assert await _count(factory, FundamentalSnapshot) == 4 + + +async def test_no_op_on_repeated_revision(engine): + factory = _factory(engine) + await run_import(FakeImporter("rev1", n_rows=4), engine=engine) + + second = FakeImporter("rev1", n_rows=4) + run = await run_import(second, engine=engine) + + assert run is not None and run.status == STATUS_NO_OP + assert second.staged_called is False # never fetched + assert await _count(factory, FundamentalSnapshot) == 4 # unchanged + + runs = await _runs(factory) + assert [r.status for r in runs] == [STATUS_PROMOTED, STATUS_NO_OP] + + +async def test_new_revision_after_promote_stages_again(engine): + factory = _factory(engine) + await run_import(FakeImporter("rev1", n_rows=2), engine=engine) + run = await run_import(FakeImporter("rev2", n_rows=3), engine=engine) + + assert run is not None and run.status == STATUS_PROMOTED + assert await _count(factory, FundamentalSnapshot) == 5 # 2 + 3 + + +async def test_failed_validation_leaves_data_untouched(engine): + factory = _factory(engine) + await run_import(FakeImporter("rev1", n_rows=3), engine=engine) # baseline + + run = await run_import(FakeImporter("rev2", ok=False, n_rows=5), engine=engine) + + assert run is not None and run.status == STATUS_FAILED + assert "coverage" in (run.error_details or "") + assert await _count(factory, FundamentalSnapshot) == 3 # untouched + assert await _count(factory, SystemEvent) == 1 # alerted + + +async def test_exception_in_promote_rolls_back(engine): + factory = _factory(engine) + await run_import(FakeImporter("rev1", n_rows=3), engine=engine) # baseline + + run = await run_import(FakeImporter("rev2", raise_in="promote"), engine=engine) + + assert run is not None and run.status == STATUS_FAILED + assert "boom-promote" in (run.error_details or "") + assert await _count(factory, FundamentalSnapshot) == 3 # partial write rolled back + assert await _count(factory, SystemEvent) == 1