"""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, timedelta, 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_DEFERRED = "deferred" 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) # Expected source-side lag: retry without an immediate error alert. Sources # can bound the quiet period with deferred_alert_after_days. Only meaningful # when ok=False. retryable: bool = False deferred_alert_after_days: int | None = None @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, run_id: int) -> dict[str, int]: """Apply ``staged`` to the live tables. Called inside the promotion transaction; the caller commits. ``run_id`` is the current ``data_import_runs.id`` so written rows can be stamped with their ``import_run_id``. 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() async def _has_promoted_since(db: AsyncSession, source: str, cutoff: datetime) -> bool: row = await db.execute( select(DataImportRun.id) .where( DataImportRun.source == source, DataImportRun.status == STATUS_PROMOTED, DataImportRun.started_at >= cutoff, ) .limit(1) ) return row.scalar_one_or_none() is not None def _now() -> datetime: return datetime.now(timezone.utc) async def _alert( db: AsyncSession, source: str, code: str, messages: list[str], *, severity: str = "error", dedup_hours: int = 24, ) -> None: try: await system_event_service.log_event( db, severity=severity, source="data_import", code=f"{source}_{code}", message=(("; ".join(messages)) or code)[:_MAX_ERROR_LEN], dedup_key=f"data_import:{source}:{code}", dedup_hours=dedup_hours, ) 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, force: bool = False, ) -> DataImportRun | None: """Run one import for ``importer``. Returns the recorded ``DataImportRun`` (promoted / no_op / deferred / failed), or None when the per-source advisory lock is already held (another run is active). ``force`` runs even when the revision is unchanged. The revision tracks the *source*, so a re-import driven by a change on our side — a parser fix that makes stored rows stale — is a no_op under the normal gate. Manually invoked only; scheduled jobs must leave it False so an unchanged source stays a no_op. """ 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 not force and 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.error_details = ("; ".join(result.messages))[:_MAX_ERROR_LEN] run.completed_at = _now() if result.retryable: run.status = STATUS_DEFERRED await session.commit() alert_days = result.deferred_alert_after_days if alert_days is not None: alert_days = max(1, alert_days) cutoff = run.started_at - timedelta(days=alert_days) if not await _has_promoted_since(session, source, cutoff): await _alert( session, source, "deferred_stale", [ f"No successful promotion for at least " f"{alert_days} day(s); deferred source lag may " f"now hide aged-out unresolved items: " f"{run.error_details or 'validation deferred'}" ], severity="warning", dedup_hours=alert_days * 24, ) logger.info( "data_import %s: deferred for retry: %s", source, result.messages, ) return run run.status = STATUS_FAILED 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.id) 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()