"""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 asyncio 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 # Record the attempt FIRST — before the external revision probe, the # most likely failure — so anything below is recorded and alerted and # never escapes unrecorded. Revision is filled in once detected. run = DataImportRun( source=source, status=STATUS_RUNNING, started_at=_now(), ) session.add(run) await session.commit() await session.refresh(run) try: revision = await importer.detect_revision(session) run.revision = revision last_rev = await _last_promoted_revision(session, source) if revision is not None and revision == last_rev: run.status = STATUS_NO_OP run.completed_at = _now() await session.commit() logger.info("data_import %s: no_op (revision %s)", source, revision) return run 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 asyncio.CancelledError: # Deploy / scheduler shutdown: best-effort mark failed so no # ``running`` row lingers, then let the cancellation propagate — # never swallow it. try: await session.rollback() run.status = STATUS_FAILED run.error_details = "cancelled" run.completed_at = _now() await session.commit() except BaseException: # noqa: BLE001 — best-effort during teardown logger.warning( "data_import %s: could not record cancellation", source ) raise 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()