diff --git a/app/services/data_import.py b/app/services/data_import.py index 7fbfb5e..11f1adf 100644 --- a/app/services/data_import.py +++ b/app/services/data_import.py @@ -94,9 +94,11 @@ class SourceImporter(Protocol): """Run the source's validation gates against ``staged``. Read-only.""" ... - async def promote(self, db: AsyncSession, staged: Any) -> dict[str, int]: + 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. Returns row-count deltas.""" + 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.""" ... @@ -213,7 +215,7 @@ async def run_import( return run # Promotion: importer writes + run-row flip in one transaction. - row_counts = await importer.promote(session, staged) + 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() diff --git a/tests/unit/test_data_import_framework.py b/tests/unit/test_data_import_framework.py index 67b200c..02a60a8 100644 --- a/tests/unit/test_data_import_framework.py +++ b/tests/unit/test_data_import_framework.py @@ -96,18 +96,19 @@ class FakeImporter: messages=[] if self.ok else ["coverage below threshold"], ) - async def promote(self, db, staged): + async def promote(self, db, staged, run_id): 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)) + db.add(_snapshot(self.revision, i, run_id=run_id)) self.promoted = True + self.promoted_run_id = run_id return {"fundamental_snapshots": len(staged)} -def _snapshot(revision: str, i: int) -> FundamentalSnapshot: +def _snapshot(revision: str, i: int, run_id: int | None = None) -> FundamentalSnapshot: return FundamentalSnapshot( cik=f"{i:010d}", accession=f"{revision}-{i:06d}", @@ -118,6 +119,7 @@ def _snapshot(revision: str, i: int) -> FundamentalSnapshot: fiscal_year=2026, fiscal_period="Q2", revenue=1000.0 + i, + import_run_id=run_id, ) @@ -154,6 +156,12 @@ async def test_promote_writes_and_records_run(engine): assert run.source_max_date == date(2026, 7, 21) assert run.completed_at is not None assert await _count(factory, FundamentalSnapshot) == 4 + # rows stamped with the run id + async with factory() as s: + stamped = ( + await s.execute(select(FundamentalSnapshot.import_run_id)) + ).scalars().all() + assert stamped and all(rid == run.id for rid in stamped) async def test_no_op_on_repeated_revision(engine):