feat(dolt): source-agnostic import-run framework (A1)
run_import + SourceImporter Protocol (detect_revision/stage/validate/promote) giving every bulk importer the plan's non-negotiables, KISS: - one run per source at a time — Postgres session-level advisory lock held on a single pinned engine.connect() so it survives the running-row and promotion commits; no-op on SQLite. - idempotent per revision — cheap detect_revision compared to the last promoted run; unchanged revision records a no_op with zero writes (no fetch). - staging (in-memory, no physical staging tables) → validate (read-only) → atomic promote + run-row flip in one transaction. - failed validation or mid-run exception marks the run failed, alerts via system_event_service, and leaves live tables untouched. Every attempt recorded in data_import_runs; conflicts summary in validation_json (no conflicts table). Concrete SEC/earnings importers land in later phases. Tests: 6 orchestration tests (no_op / promote / new-revision / failed-untouched / promote-exception-rollback) + deterministic advisory-key derivation. Full suite 680 passed. Advisory-lock mutual exclusion is PG-verify-pending (SQLite no-ops it — flagged, not covered). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Claude Opus 4.8
@@ -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()
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user