from __future__ import annotations import unittest from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker from govoplan_core.auth import ApiPrincipal from govoplan_core.core.access import PrincipalRef from govoplan_core.core.datasources import ( CAPABILITY_DATASOURCE_ORIGINS, DatasourceAccessError, DatasourceField, DatasourceOrigin, DatasourceOriginReadRequest, DatasourceOriginReadResult, DatasourcePublicationRequest, DatasourceReadRequest, DatasourceStageInput, DatasourceValidationError, ) from govoplan_core.db.base import Base, utcnow from govoplan_datasources.backend.db.models import ( DatasourceMaterializationRecord, DatasourcePublicationRecord, DatasourceRecord, DatasourceStageRecord, ) from govoplan_datasources.backend.service import ( CATALOGUE_READ_SCOPE, SOURCE_WRITE_SCOPE, STAGE_WRITE_SCOPE, SqlDatasourceProvider, ) def principal( tenant_id: str = "tenant-1", *, scopes: tuple[str, ...] = ( CATALOGUE_READ_SCOPE, SOURCE_WRITE_SCOPE, STAGE_WRITE_SCOPE, ), ) -> ApiPrincipal: return ApiPrincipal( principal=PrincipalRef( account_id="account-1", membership_id="membership-1", tenant_id=tenant_id, scopes=frozenset(scopes), ), account=object(), user=object(), ) class FakeOriginProvider: def __init__(self) -> None: self.rows = [{"id": 1, "name": "Initial"}] def _origin(self) -> DatasourceOrigin: return DatasourceOrigin( ref="snapshot:origin-1", source_name="connector_cases", name="Connector cases", kind="database", shape="tabular", supported_modes=("live", "cached"), provider="connectors.test", schema=( DatasourceField(name="id", data_type="integer", nullable=False), DatasourceField(name="name", data_type="string", nullable=False), ), schema_version="1", fingerprint=f"version-{len(self.rows)}-{self.rows[-1]['name']}", row_count=len(self.rows), updated_at=utcnow(), ) def list_origins( self, _session: object, _principal: object, *, query: str = "", limit: int = 100, ): origin = self._origin() return (origin,) if query.lower() in origin.name.lower() else () def get_origin( self, _session: object, _principal: object, *, origin_ref: str, ): return self._origin() if origin_ref == "snapshot:origin-1" else None def read_origin( self, _session: object, _principal: object, *, request: DatasourceOriginReadRequest, ) -> DatasourceOriginReadResult: origin = self.get_origin(None, None, origin_ref=request.origin_ref) if origin is None: raise AssertionError("Unexpected origin") rows = self.rows[request.offset : request.offset + request.limit] return DatasourceOriginReadResult( origin=origin, rows=tuple(dict(row) for row in rows), total_rows=len(self.rows), truncated=request.offset + len(rows) < len(self.rows), ) class FakeRegistry: def __init__(self, origin_provider: FakeOriginProvider) -> None: self.origin_provider = origin_provider def has_capability(self, name: str) -> bool: return name == CAPABILITY_DATASOURCE_ORIGINS def capability(self, name: str) -> object: if not self.has_capability(name): raise KeyError(name) return self.origin_provider class DatasourceLifecycleTests(unittest.TestCase): def setUp(self) -> None: self.engine = create_engine("sqlite:///:memory:") Base.metadata.create_all( self.engine, tables=[ DatasourceRecord.__table__, DatasourceMaterializationRecord.__table__, DatasourceStageRecord.__table__, DatasourcePublicationRecord.__table__, ], ) self.Session = sessionmaker(bind=self.engine) self.session = self.Session() self.origins = FakeOriginProvider() self.provider = SqlDatasourceProvider( registry=FakeRegistry(self.origins), ) def tearDown(self) -> None: self.session.close() Base.metadata.drop_all( self.engine, tables=[ DatasourceStageRecord.__table__, DatasourcePublicationRecord.__table__, DatasourceMaterializationRecord.__table__, DatasourceRecord.__table__, ], ) self.engine.dispose() def test_static_stage_promote_update_and_frozen_read(self) -> None: first_stage = self.provider.create_stage( self.session, principal(), stage=DatasourceStageInput( name="Monthly cases", source_name="monthly_cases", kind="upload", mode="static", shape="tabular", rows=({"id": 1, "status": "new"},), ), ) datasource, first = self.provider.promote_stage( self.session, principal(), stage_ref=first_stage.ref, ) frozen = self.provider.freeze_datasource( self.session, principal(), datasource_ref=datasource.ref, label="Import evidence", ) second_stage = self.provider.create_stage( self.session, principal(), stage=DatasourceStageInput( name="Monthly cases", source_name="monthly_cases", kind="upload", mode="static", shape="tabular", target_datasource_ref=datasource.ref, rows=({"id": 2, "status": "checked", "note": "schema changed"},), ), ) updated, second = self.provider.promote_stage( self.session, principal(), stage_ref=second_stage.ref, ) self.session.commit() current = self.provider.read_datasource( self.session, principal(), request=DatasourceReadRequest(datasource_ref=datasource.ref), ) frozen_result = self.provider.read_datasource( self.session, principal(), request=DatasourceReadRequest( datasource_ref=datasource.ref, consistency="frozen", ), ) self.assertEqual(1, first.revision) self.assertEqual(2, frozen.revision) self.assertEqual(3, second.revision) self.assertEqual( [{"id": 2, "status": "checked", "note": "schema changed"}], list(current.rows), ) self.assertEqual([{"id": 1, "status": "new"}], list(frozen_result.rows)) self.assertEqual(second.ref, updated.current_materialization_ref) self.assertEqual("2", updated.schema_version) self.assertEqual( "2", current.datasource.schema_version, ) self.assertEqual(frozen.ref, frozen_result.materialization.ref) def test_live_reads_origin_and_cached_refresh_is_explicit(self) -> None: live = self.provider.register_origin( self.session, principal(), origin_ref="snapshot:origin-1", name="Live cases", source_name="live_cases", mode="live", ) cached = self.provider.register_origin( self.session, principal(), origin_ref="snapshot:origin-1", name="Cached cases", source_name="cached_cases", mode="cached", ) self.session.commit() self.origins.rows = [ {"id": 1, "name": "Initial"}, {"id": 2, "name": "New"}, ] live_read = self.provider.read_datasource( self.session, principal(), request=DatasourceReadRequest(datasource_ref=live.ref), ) cached_before = self.provider.read_datasource( self.session, principal(), request=DatasourceReadRequest(datasource_ref=cached.ref), ) cached_live = self.provider.read_datasource( self.session, principal(), request=DatasourceReadRequest( datasource_ref=cached.ref, consistency="live", ), ) refreshed, materialization = self.provider.refresh_datasource( self.session, principal(), datasource_ref=cached.ref, ) cached_after = self.provider.read_datasource( self.session, principal(), request=DatasourceReadRequest(datasource_ref=cached.ref), ) self.assertEqual(2, live_read.total_rows) self.assertEqual(1, cached_before.total_rows) self.assertEqual(2, cached_live.total_rows) self.assertEqual(2, cached_after.total_rows) self.assertEqual(materialization.ref, refreshed.current_materialization_ref) def test_tenant_and_scope_isolation(self) -> None: stage = self.provider.create_stage( self.session, principal(), stage=DatasourceStageInput( name="Private", source_name="private", kind="upload", mode="static", shape="tabular", rows=({"id": 1},), ), ) datasource, _ = self.provider.promote_stage( self.session, principal(), stage_ref=stage.ref, ) self.session.commit() self.assertEqual( (), self.provider.list_datasources(self.session, principal("tenant-2")), ) self.assertIsNone( self.provider.get_datasource( self.session, principal("tenant-2"), datasource_ref=datasource.ref, ) ) with self.assertRaises(DatasourceAccessError): self.provider.list_datasources( self.session, principal(scopes=()), ) def test_stage_writer_can_list_stages_without_catalogue_scope(self) -> None: writer = principal(scopes=(STAGE_WRITE_SCOPE,)) stage = self.provider.create_stage( self.session, writer, stage=DatasourceStageInput( name="Pending import", source_name="pending_import", kind="upload", mode="static", shape="tabular", rows=({"id": 1},), ), ) self.assertEqual( (stage.ref,), tuple(item.ref for item in self.provider.list_stages(self.session, writer)), ) def test_producer_publication_is_atomic_idempotent_and_addressable(self) -> None: producer = principal(scopes=(SOURCE_WRITE_SCOPE,)) request = DatasourcePublicationRequest( producer_module="dataflow", producer_run_ref="dataflow-run:run-1", idempotency_key="monthly-output-2026-07", name="Monthly comparison result", source_name="monthly_comparison_result", rows=( {"case_id": "A-1", "result": "match"}, {"case_id": "A-2", "result": "review"}, ), freeze=True, frozen_label="July 2026", provenance={"pipeline_revision": 3}, ) first = self.provider.publish_rows( self.session, producer, request=request, ) replay = self.provider.publish_rows( self.session, producer, request=request, ) second = self.provider.publish_rows( self.session, producer, request=DatasourcePublicationRequest( producer_module="dataflow", producer_run_ref="dataflow-run:run-2", idempotency_key="monthly-output-2026-08", target_datasource_ref=first.datasource.ref, rows=({"case_id": "A-3", "result": "match"},), ), ) self.session.commit() self.assertFalse(first.replayed) self.assertTrue(replay.replayed) self.assertEqual(first.ref, replay.ref) self.assertEqual(first.materialization.ref, replay.materialization.ref) self.assertIsNotNone(first.materialization.frozen_at) self.assertEqual(1, first.materialization.revision) self.assertEqual(2, second.materialization.revision) self.assertEqual(first.datasource.ref, second.datasource.ref) self.assertEqual( 2, self.session.query(DatasourcePublicationRecord).count(), ) def test_publication_idempotency_key_rejects_different_output(self) -> None: producer = principal(scopes=(SOURCE_WRITE_SCOPE,)) base = DatasourcePublicationRequest( producer_module="dataflow", producer_run_ref="dataflow-run:run-1", idempotency_key="stable-key", name="Result", source_name="result", rows=({"id": 1},), ) self.provider.publish_rows(self.session, producer, request=base) with self.assertRaises(DatasourceValidationError): self.provider.publish_rows( self.session, producer, request=DatasourcePublicationRequest( producer_module="dataflow", producer_run_ref="dataflow-run:run-1", idempotency_key="stable-key", name="Result", source_name="result", rows=({"id": 2},), ), ) def test_publication_requires_source_write_scope(self) -> None: with self.assertRaises(DatasourceAccessError): self.provider.publish_rows( self.session, principal(scopes=()), request=DatasourcePublicationRequest( producer_module="dataflow", producer_run_ref="dataflow-run:run-1", idempotency_key="denied", name="Result", source_name="result", rows=({"id": 1},), ), ) if __name__ == "__main__": unittest.main()