349 lines
14 KiB
Python
349 lines
14 KiB
Python
"""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 exists, 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
|
|
deferred_alert_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, 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 _promotion_state_since(
|
|
db: AsyncSession, source: str, cutoff: datetime
|
|
) -> str:
|
|
promoted = (
|
|
DataImportRun.source == source,
|
|
DataImportRun.status == STATUS_PROMOTED,
|
|
)
|
|
ever, recent = (
|
|
await db.execute(
|
|
select(
|
|
exists().where(*promoted),
|
|
exists().where(*promoted, DataImportRun.started_at >= cutoff),
|
|
)
|
|
)
|
|
).one()
|
|
return "recent" if recent else "stale" if ever else "never"
|
|
|
|
|
|
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)
|
|
promotion_state = await _promotion_state_since(
|
|
session, source, cutoff
|
|
)
|
|
if promotion_state != "recent":
|
|
history = (
|
|
f"{source} import has never promoted successfully"
|
|
if promotion_state == "never"
|
|
else f"{source} import has not promoted successfully "
|
|
f"within {alert_days} day(s)"
|
|
)
|
|
await _alert(
|
|
session,
|
|
source,
|
|
"deferred_stale",
|
|
[
|
|
f"{history}; import remains deferred",
|
|
*result.deferred_alert_messages,
|
|
f"Current deferral: "
|
|
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()
|