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.change_sequence import ChangeSequenceEntry from govoplan_core.core.datasources import ( CAPABILITY_DATASOURCE_ORIGINS, DatasourceAccessError, DatasourceArtifactReference, DatasourceField, DatasourceGovernance, DatasourceOrigin, DatasourceOriginReadRequest, DatasourceOriginReadResult, DatasourcePublicationRequest, DatasourceReadRequest, DatasourceStageInput, DatasourceUnavailableError, DatasourceValidationError, ) from govoplan_core.core.tabular_sources import ( TabularPreviewDiagnostic, TabularPushdown, TabularSourceHealth, ) from govoplan_core.db.base import Base, utcnow from govoplan_datasources.backend.db.models import ( DatasourceGovernanceReferenceRecord, DatasourceMaterializationRecord, DatasourcePayloadRecord, DatasourcePayloadRowRecord, DatasourcePublicationRecord, DatasourceRecord, DatasourceStageRecord, ) from govoplan_datasources.backend.service import ( CATALOGUE_READ_SCOPE, SOURCE_WRITE_SCOPE, STAGE_WRITE_SCOPE, SqlDatasourceProvider, ) from govoplan_datasources.backend.payloads import ( ExternalArtifactPayloadBackend, create_database_rows_payload, finalize_payload_deletion, mark_unreferenced_payload_for_deletion, verify_payload_integrity, ) 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(), source_mode="cached", pushdown=TabularPushdown(projections=True, pagination=True), health=TabularSourceHealth( status="healthy", code="snapshot.ready", summary="The immutable snapshot is ready.", ), ) 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), returned_bytes=64, elapsed_ms=4, effective_row_limit=request.limit, effective_byte_limit=request.max_bytes, effective_timeout_ms=request.timeout_ms, diagnostics=( TabularPreviewDiagnostic( severity="info", code="preview.complete", message="The bounded preview completed.", ), ), ) 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 FakeArtifactBackend: backend = "test_artifact" def __init__(self) -> None: self.verified: list[str] = [] self.deleted: list[str] = [] def read_rows( self, _session, *, tenant_id: str, artifact: DatasourceArtifactReference, offset: int, limit: int, ): self.assert_tenant(tenant_id) stop = min(artifact.row_count, offset + limit) return tuple( {"id": index, "result": "match"} for index in range(offset, stop) ) def verify( self, _session, *, tenant_id: str, artifact: DatasourceArtifactReference, ) -> None: self.assert_tenant(tenant_id) if not artifact.locator.startswith("artifact:"): raise DatasourceUnavailableError("Unknown test artifact.") self.verified.append(artifact.locator) def delete( self, _session, *, tenant_id: str, artifact: DatasourceArtifactReference, ) -> None: self.assert_tenant(tenant_id) self.deleted.append(artifact.locator) def assert_tenant(self, tenant_id: str) -> None: if tenant_id != "tenant-1": raise AssertionError("Artifact backend crossed a tenant boundary.") class DatasourceLifecycleTests(unittest.TestCase): def setUp(self) -> None: self.engine = create_engine("sqlite:///:memory:") Base.metadata.create_all( self.engine, tables=[ DatasourceRecord.__table__, DatasourceGovernanceReferenceRecord.__table__, DatasourcePayloadRecord.__table__, DatasourcePayloadRowRecord.__table__, DatasourceMaterializationRecord.__table__, DatasourceStageRecord.__table__, DatasourcePublicationRecord.__table__, ChangeSequenceEntry.__table__, ], ) self.Session = sessionmaker(bind=self.engine) self.session = self.Session() self.origins = FakeOriginProvider() self.artifacts = FakeArtifactBackend() self.provider = SqlDatasourceProvider( registry=FakeRegistry(self.origins), payload_backends=(ExternalArtifactPayloadBackend(self.artifacts),), ) def tearDown(self) -> None: self.session.close() Base.metadata.drop_all( self.engine, tables=[ DatasourceStageRecord.__table__, DatasourcePublicationRecord.__table__, DatasourceMaterializationRecord.__table__, DatasourcePayloadRowRecord.__table__, DatasourcePayloadRecord.__table__, DatasourceGovernanceReferenceRecord.__table__, DatasourceRecord.__table__, ChangeSequenceEntry.__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) first_record = self.session.get( DatasourceMaterializationRecord, first.ref.removeprefix("materialization:"), ) frozen_record = self.session.get( DatasourceMaterializationRecord, frozen.ref.removeprefix("materialization:"), ) self.assertIsNotNone(first_record) self.assertIsNotNone(frozen_record) self.assertEqual(first_record.payload_id, frozen_record.payload_id) self.assertEqual([], first_record.rows) def test_stage_quality_and_schema_gates_block_only_error_diagnostics(self) -> None: governance = DatasourceGovernance( authority_mode="native_authoritative", publication_state="internal", quality_policy={ "version": "monthly-cases-v1", "rules": [ {"id": "unique-id", "type": "unique", "fields": ["id"]}, ], "schema_policy": {"field_added_required": "warning"}, }, ) first_stage = self.provider.create_stage( self.session, principal(), stage=DatasourceStageInput( name="Monthly cases", source_name="monthly_quality_cases", kind="upload", mode="static", shape="tabular", rows=({"id": 1},), governance=governance, ), ) datasource, _first = self.provider.promote_stage( self.session, principal(), stage_ref=first_stage.ref, ) blocked_stage = self.provider.create_stage( self.session, principal(), stage=DatasourceStageInput( name="Monthly cases", source_name="monthly_quality_cases", kind="upload", mode="static", shape="tabular", target_datasource_ref=datasource.ref, rows=({"id": 2, "name": "Ada"}, {"id": 2, "name": "Lin"}), ), ) self.assertEqual("invalid", blocked_stage.state) self.assertFalse(blocked_stage.validation["valid"]) self.assertEqual( "quality.unique", blocked_stage.validation["errors"][0]["code"], ) self.assertEqual( "warning", blocked_stage.validation["schema_change"]["classification"], ) with self.assertRaisesRegex( DatasourceValidationError, "blocking quality or schema diagnostics", ): self.provider.promote_stage( self.session, principal(), stage_ref=blocked_stage.ref, ) ready_stage = self.provider.create_stage( self.session, principal(), stage=DatasourceStageInput( name="Monthly cases", source_name="monthly_quality_cases", kind="upload", mode="static", shape="tabular", target_datasource_ref=datasource.ref, rows=({"id": 2, "name": "Ada"}, {"id": 3, "name": "Lin"}), ), ) self.assertEqual("ready", ready_stage.state) self.assertTrue(ready_stage.validation["valid"]) self.assertEqual("schema.field_added_required", ready_stage.validation["warnings"][0]["code"]) _updated, materialization = self.provider.promote_stage( self.session, principal(), stage_ref=ready_stage.ref, ) self.assertEqual( ready_stage.validation["policy_hash"], materialization.provenance["stage_validation"]["policy_hash"], ) def test_governance_is_queryable_and_snapshotted_per_materialization(self) -> None: stage = self.provider.create_stage( self.session, principal(), stage=DatasourceStageInput( name="Governed register", source_name="governed_register", kind="upload", mode="static", shape="tabular", rows=({"id": 1},), governance=DatasourceGovernance( owner_ref="function:data-owner", steward_ref="account:steward", responsible_organization_ref="organization:office-1", authority_mode="native_authoritative", legal_basis_refs=("policy:register-use",), purposes=("case_processing",), semantic_definition="Authoritative case register export.", official_keys=("id",), classification="restricted", publication_state="internal", quality_policy={"required_keys": ["id"]}, affected_refs=("service:permit", "report:monthly"), dependency_refs=("dataflow:monthly-case-check",), ), ), ) datasource, first = self.provider.promote_stage( self.session, principal(), stage_ref=stage.ref, ) self.assertEqual("function:data-owner", datasource.governance.owner_ref) self.assertEqual("restricted", first.governance.classification) self.assertEqual( [datasource.ref], [ item.ref for item in self.provider.list_datasources( self.session, principal(), authority_mode="native_authoritative", classification="restricted", publication_state="internal", owner_ref="function:data-owner", responsible_organization_ref="organization:office-1", ) ], ) self.assertEqual( [datasource.ref], [ item.ref for item in self.provider.list_datasources( self.session, principal(), affected_ref="service:permit", dependency_ref="dataflow:monthly-case-check", ) ], ) self.assertEqual( (), self.provider.list_datasources( self.session, principal("tenant-2"), affected_ref="service:permit", ), ) changed = self.provider.update_datasource_governance( self.session, principal(), datasource_ref=datasource.ref, governance=DatasourceGovernance( owner_ref="function:new-owner", authority_mode="native_authoritative", purposes=("case_processing",), classification="confidential", publication_state="internal", ), ) history = self.provider.list_materializations( self.session, principal(), datasource_ref=datasource.ref, ) self.assertEqual("function:new-owner", changed.governance.owner_ref) self.assertEqual("function:data-owner", history[0].governance.owner_ref) self.assertEqual("restricted", history[0].governance.classification) self.assertEqual( (), self.provider.list_datasources( self.session, principal(), dependency_ref="dataflow:monthly-case-check", ), ) 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(64, live_read.returned_bytes) self.assertEqual("preview.complete", live_read.diagnostics[0].code) self.assertEqual( "cached", live_read.datasource.metadata["source_contract"]["source_mode"], ) self.assertTrue( live_read.datasource.metadata["source_contract"]["pushdown"][ "projections" ] ) 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) self.assertEqual( "healthy", refreshed.metadata["source_contract"]["health"]["status"], ) 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_artifact_publication_pins_large_payload_and_supports_bounded_reads( self, ) -> None: artifact = DatasourceArtifactReference( backend="test_artifact", locator="artifact:monthly-output", checksum="a" * 64, row_count=25_000, byte_count=12_000_000, schema=( DatasourceField("id", "integer", nullable=False), DatasourceField("result", "string", nullable=False), ), fingerprint="b" * 64, ) published = self.provider.publish_rows( self.session, principal(scopes=(SOURCE_WRITE_SCOPE, CATALOGUE_READ_SCOPE)), request=DatasourcePublicationRequest( producer_module="dataflow", producer_run_ref="dataflow-run:large-output", idempotency_key="large-output", name="Large output", source_name="large_output", artifact=artifact, ), ) preview = self.provider.read_datasource( self.session, principal(scopes=(CATALOGUE_READ_SCOPE,)), request=DatasourceReadRequest( datasource_ref=published.datasource.ref, offset=10, limit=3, ), ) self.assertEqual("published", published.status) self.assertEqual(25_000, published.materialization.row_count) self.assertEqual( [{"id": 10, "result": "match"}, {"id": 11, "result": "match"}, {"id": 12, "result": "match"}], list(preview.rows), ) self.assertEqual( ["artifact:monthly-output", "artifact:monthly-output"], self.artifacts.verified, ) def test_unattested_artifact_quality_rules_require_review_without_becoming_current( self, ) -> None: published = self.provider.publish_rows( self.session, principal(scopes=(SOURCE_WRITE_SCOPE,)), request=DatasourcePublicationRequest( producer_module="reporting", producer_run_ref="report-run:review", idempotency_key="review-output", name="Review output", source_name="review_output", artifact=DatasourceArtifactReference( backend="test_artifact", locator="artifact:review-output", checksum="c" * 64, row_count=2, byte_count=128, schema=(DatasourceField("id", "integer", False),), fingerprint="d" * 64, ), governance=DatasourceGovernance( quality_policy={ "version": "unique-id-v1", "rules": [ { "id": "unique-id", "type": "unique", "fields": ["id"], } ], } ), ), ) record = self.session.get( DatasourceRecord, published.datasource.ref.removeprefix("datasource:"), ) self.assertEqual("review_required", published.status) self.assertEqual("review_required", published.materialization.state) self.assertIsNotNone(record) assert record is not None self.assertIsNone(record.current_materialization_id) def test_artifact_warning_is_a_notification_ready_terminal_state(self) -> None: published = self.provider.publish_rows( self.session, principal(scopes=(SOURCE_WRITE_SCOPE,)), request=DatasourcePublicationRequest( producer_module="dataflow", producer_run_ref="dataflow-run:warning", idempotency_key="warning-output", name="Warning output", source_name="warning_output", artifact=DatasourceArtifactReference( backend="test_artifact", locator="artifact:warning-output", checksum="e" * 64, row_count=1, byte_count=64, schema=(DatasourceField("id", "integer", False),), fingerprint="f" * 64, validation={ "status": "warning", "warnings": [ { "severity": "warning", "code": "producer.partial_match", "message": "One source used a fallback match.", } ], }, ), ), ) record = self.session.get( DatasourcePublicationRecord, published.ref.removeprefix("publication:"), ) self.assertEqual("published_with_warnings", published.status) self.assertIsNotNone(record) assert record is not None self.assertEqual( "published_with_warnings", record.status, ) 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_enforces_quality_before_persisting_any_effect(self) -> None: producer = principal(scopes=(SOURCE_WRITE_SCOPE,)) blocked = DatasourcePublicationRequest( producer_module="dataflow", producer_run_ref="dataflow-run:blocked", idempotency_key="blocked-output", name="Blocked output", source_name="blocked_output", rows=({"id": 1},), governance=DatasourceGovernance( quality_policy={ "version": "minimum-two-v1", "rules": [ { "id": "minimum-two", "type": "row_count", "minimum": 2, } ], }, ), ) with self.assertRaisesRegex( DatasourceValidationError, r"quality or schema validation \(quality\.row_count\)", ): self.provider.publish_rows(self.session, producer, request=blocked) self.assertEqual(0, self.session.query(DatasourceRecord).count()) self.assertEqual( 0, self.session.query(DatasourceMaterializationRecord).count(), ) self.assertEqual( 0, self.session.query(DatasourcePublicationRecord).count(), ) def test_publication_preserves_quality_and_schema_evidence(self) -> None: producer = principal(scopes=(SOURCE_WRITE_SCOPE,)) first = self.provider.publish_rows( self.session, producer, request=DatasourcePublicationRequest( producer_module="dataflow", producer_run_ref="dataflow-run:first", idempotency_key="governed-output-first", name="Governed output", source_name="governed_output", rows=({"id": 1},), governance=DatasourceGovernance( quality_policy={ "version": "governed-output-v1", "rules": [ { "id": "unique-id", "type": "unique", "fields": ["id"], }, ], "schema_policy": {"field_added_required": "warning"}, }, ), ), ) second = self.provider.publish_rows( self.session, producer, request=DatasourcePublicationRequest( producer_module="dataflow", producer_run_ref="dataflow-run:second", idempotency_key="governed-output-second", target_datasource_ref=first.datasource.ref, rows=({"id": 2, "result": "match"},), ), ) validation = second.materialization.provenance["publication_validation"] self.assertTrue(validation["valid"]) self.assertEqual("governed-output-v1", validation["policy_version"]) self.assertEqual( "warning", validation["schema_change"]["classification"], ) self.assertEqual( "schema.field_added_required", validation["warnings"][0]["code"], ) record = self.session.get( DatasourcePublicationRecord, second.ref.removeprefix("publication:"), ) self.assertIsNotNone(record) assert record is not None self.assertEqual( validation["policy_hash"], record.details_["validation"]["policy_hash"], ) 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},), ), ) def test_payload_preview_is_paged_and_metadata_mismatch_is_rejected(self) -> None: stage = self.provider.create_stage( self.session, principal(), stage=DatasourceStageInput( name="Paged", source_name="paged", kind="upload", mode="static", shape="tabular", rows=tuple({"id": index} for index in range(20)), ), ) datasource, materialization = self.provider.promote_stage( self.session, principal(), stage_ref=stage.ref, ) self.session.commit() result = self.provider.read_datasource( self.session, principal(), request=DatasourceReadRequest( datasource_ref=datasource.ref, offset=7, limit=3, ), ) self.assertEqual([{"id": 7}, {"id": 8}, {"id": 9}], list(result.rows)) record = self.session.get( DatasourceMaterializationRecord, materialization.ref.removeprefix("materialization:"), ) self.assertIsNotNone(record) self.assertEqual([], record.rows) self.assertEqual( 20, self.session.query(DatasourcePayloadRowRecord) .filter(DatasourcePayloadRowRecord.payload_id == record.payload_id) .count(), ) payload = self.session.get(DatasourcePayloadRecord, record.payload_id) self.assertIsNotNone(payload) verify_payload_integrity(self.session, payload) payload.row_count += 1 self.session.flush() with self.assertRaises(DatasourceUnavailableError): self.provider.read_datasource( self.session, principal(), request=DatasourceReadRequest(datasource_ref=datasource.ref), ) def test_payload_deletion_is_staged_and_reference_safe(self) -> None: payload = create_database_rows_payload( self.session, tenant_id="tenant-1", rows=({"id": 1}, {"id": 2}), actor_id="account-1", ) self.session.commit() self.assertTrue( mark_unreferenced_payload_for_deletion(self.session, payload) ) self.assertEqual("deleting", payload.state) self.session.commit() finalize_payload_deletion(self.session, payload) self.session.commit() self.assertEqual(0, self.session.query(DatasourcePayloadRecord).count()) self.assertEqual( 0, self.session.query(DatasourcePayloadRowRecord).count(), ) def test_rolled_back_materialization_leaves_no_payload_rows(self) -> None: stage = self.provider.create_stage( self.session, principal(), stage=DatasourceStageInput( name="Rollback", source_name="rollback", kind="upload", mode="static", shape="tabular", rows=({"id": 1},), ), ) self.provider.promote_stage( self.session, principal(), stage_ref=stage.ref, ) self.session.rollback() self.assertEqual(0, self.session.query(DatasourcePayloadRecord).count()) self.assertEqual( 0, self.session.query(DatasourcePayloadRowRecord).count(), ) self.assertEqual( 0, self.session.query(DatasourceMaterializationRecord).count(), ) if __name__ == "__main__": unittest.main()