from __future__ import annotations import unittest from datetime import timedelta from sqlalchemy import create_engine, select 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, DatasourceField, DatasourceGovernance, DatasourceOrigin, DatasourceOriginReadRequest, DatasourceOriginReadResult, DatasourceReadRequest, DatasourceStageInput, DatasourceUnavailableError, DatasourceValidationError, ) from govoplan_core.core.tabular_sources import TabularPushdown, TabularSourceHealth from govoplan_core.db.base import Base from govoplan_datasources.backend.db.models import ( DatasourceGovernanceReferenceRecord, DatasourceLifecycleEvidenceRecord, DatasourceMaterializationRecord, DatasourcePayloadRecord, DatasourcePayloadRowRecord, DatasourcePublicationRecord, DatasourceRecord, DatasourceStageRecord, ) from govoplan_datasources.backend.service import ( CATALOGUE_READ_SCOPE, SOURCE_WRITE_SCOPE, STAGE_APPROVE_SCOPE, STAGE_WRITE_SCOPE, SqlDatasourceProvider, ) def principal( account_id: str, *, scopes: tuple[str, ...], ) -> ApiPrincipal: return ApiPrincipal( principal=PrincipalRef( account_id=account_id, membership_id=f"membership-{account_id}", tenant_id="tenant-1", scopes=frozenset(scopes), ), account=object(), user=object(), ) class _OriginProvider: def __init__(self) -> None: self.rows = [{"id": 1, "status": "initial"}] def origin(self) -> DatasourceOrigin: return DatasourceOrigin( ref="origin:cases", source_name="cached_cases", name="Cached cases", kind="database", shape="tabular", supported_modes=("cached",), provider="connectors.test", schema=( DatasourceField("id", "integer", False), DatasourceField("status", "string", False), ), schema_version="1", fingerprint=f"cases-{self.rows[-1]['status']}", row_count=len(self.rows), source_mode="cached", pushdown=TabularPushdown(pagination=True), health=TabularSourceHealth( status="healthy", code="origin.ready", summary="Origin is ready.", ), ) def list_origins(self, _session, _principal, *, query="", limit=100): del limit origin = self.origin() return (origin,) if query.casefold() in origin.name.casefold() else () def get_origin(self, _session, _principal, *, origin_ref): return self.origin() if origin_ref == "origin:cases" else None def read_origin( self, _session, _principal, *, request: DatasourceOriginReadRequest, ) -> DatasourceOriginReadResult: rows = self.rows[request.offset : request.offset + request.limit] return DatasourceOriginReadResult( origin=self.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=1, effective_row_limit=request.limit, effective_byte_limit=request.max_bytes, effective_timeout_ms=request.timeout_ms, ) class _Registry: def __init__(self, origins: _OriginProvider) -> None: self.origins = origins def has_capability(self, name: str) -> bool: return name == CAPABILITY_DATASOURCE_ORIGINS def capability(self, name: str): if not self.has_capability(name): raise KeyError(name) return self.origins class DatasourceLifecycleGovernanceTests(unittest.TestCase): def setUp(self) -> None: self.engine = create_engine("sqlite:///:memory:") self.tables = [ DatasourceRecord.__table__, DatasourceGovernanceReferenceRecord.__table__, DatasourcePayloadRecord.__table__, DatasourcePayloadRowRecord.__table__, DatasourceMaterializationRecord.__table__, DatasourceStageRecord.__table__, DatasourcePublicationRecord.__table__, DatasourceLifecycleEvidenceRecord.__table__, ] Base.metadata.create_all(self.engine, tables=self.tables) self.session = sessionmaker(bind=self.engine)() self.origins = _OriginProvider() self.provider = SqlDatasourceProvider(registry=_Registry(self.origins)) self.writer = principal( "writer", scopes=( CATALOGUE_READ_SCOPE, SOURCE_WRITE_SCOPE, STAGE_WRITE_SCOPE, ), ) def tearDown(self) -> None: self.session.close() Base.metadata.drop_all(self.engine, tables=list(reversed(self.tables))) self.engine.dispose() def test_approval_quorum_is_attributable_idempotent_and_separated(self) -> None: stage = self.provider.create_stage( self.session, self.writer, stage=DatasourceStageInput( name="Approval register", source_name="approval_register", kind="upload", mode="static", shape="tabular", rows=({"id": 1},), governance=DatasourceGovernance( approval_policy={ "version": "promotion-v2", "required": True, "required_approvals": 2, "separation_of_duties": True, } ), ), ) self.assertEqual("awaiting_approval", stage.state) policy_hash = str(stage.approval["policy_hash"]) subject_digest = str(stage.approval["subject_digest"]) creator_approver = principal( "writer", scopes=(STAGE_APPROVE_SCOPE,), ) with self.assertRaisesRegex( DatasourceValidationError, "creator cannot approve", ): self.provider.decide_stage( self.session, creator_approver, stage_ref=stage.ref, decision="approve", reason="Creator review", expected_policy_hash=policy_hash, expected_subject_digest=subject_digest, ) first, first_hash, replayed = self.provider.decide_stage( self.session, principal("approver-1", scopes=(STAGE_APPROVE_SCOPE,)), stage_ref=stage.ref, decision="approve", reason="Quality and schema evidence reviewed.", expected_policy_hash=policy_hash, expected_subject_digest=subject_digest, ) self.assertFalse(replayed) self.assertTrue(first_hash) self.assertEqual("awaiting_approval", first.state) replay, replay_hash, replayed = self.provider.decide_stage( self.session, principal("approver-1", scopes=(STAGE_APPROVE_SCOPE,)), stage_ref=stage.ref, decision="approve", reason="Quality and schema evidence reviewed.", expected_policy_hash=policy_hash, expected_subject_digest=subject_digest, ) self.assertTrue(replayed) self.assertIsNone(replay_hash) self.assertEqual(1, replay.approval["approval_count"]) approved, second_hash, replayed = self.provider.decide_stage( self.session, principal("approver-2", scopes=(STAGE_APPROVE_SCOPE,)), stage_ref=stage.ref, decision="approve", reason="Authority and promotion impact reviewed.", expected_policy_hash=policy_hash, expected_subject_digest=subject_digest, ) self.assertFalse(replayed) self.assertTrue(second_hash) self.assertEqual("ready", approved.state) datasource, materialization = self.provider.promote_stage( self.session, self.writer, stage_ref=stage.ref, ) self.assertEqual(2, materialization.provenance["stage_approval"]["approval_count"]) self.assertTrue(materialization.provenance["promotion_evidence_hash"]) self.assertEqual(materialization.ref, datasource.current_materialization_ref) evidence = tuple( reversed( self.provider.list_lifecycle_evidence( self.session, self.writer, subject_ref=stage.ref, ) ) ) self.assertEqual( ["stage.validated", "stage.approved", "stage.approved", "stage.promoted"], [item.event_type for item in evidence], ) for previous, current in zip(evidence, evidence[1:], strict=False): self.assertEqual(previous.event_hash, current.previous_event_hash) def test_approved_refresh_stage_is_required_before_current_changes(self) -> None: datasource = self.provider.register_origin( self.session, self.writer, origin_ref="origin:cases", name="Cached cases", source_name="cached_cases", mode="cached", governance=DatasourceGovernance( approval_policy={ "version": "refresh-v1", "required": True, "required_approvals": 1, } ), ) original_ref = datasource.current_materialization_ref self.origins.rows = [{"id": 1, "status": "refreshed"}] with self.assertRaisesRegex( DatasourceValidationError, "approved refresh stage", ): self.provider.refresh_datasource( self.session, self.writer, datasource_ref=datasource.ref, ) stage = self.provider.prepare_refresh( self.session, self.writer, datasource_ref=datasource.ref, ) self.assertEqual("awaiting_approval", stage.state) approved, _, _ = self.provider.decide_stage( self.session, principal("refresh-approver", scopes=(STAGE_APPROVE_SCOPE,)), stage_ref=stage.ref, decision="approve", reason="The connector delta and schema are acceptable.", expected_policy_hash=str(stage.approval["policy_hash"]), expected_subject_digest=str(stage.approval["subject_digest"]), ) self.assertEqual("ready", approved.state) refreshed, current = self.provider.promote_stage( self.session, self.writer, stage_ref=stage.ref, ) self.assertNotEqual(original_ref, refreshed.current_materialization_ref) self.assertEqual(current.ref, refreshed.current_materialization_ref) preview = self.provider.read_datasource( self.session, self.writer, request=DatasourceReadRequest(datasource_ref=datasource.ref), ) self.assertEqual("refreshed", preview.rows[0]["status"]) def test_retention_preview_blocks_current_and_held_evidence_then_purges_payload(self) -> None: governance = DatasourceGovernance( retention_policy={ "version": "records-v3", "enabled": True, "materialization_days": 1, "frozen_evidence_days": 1, } ) first_stage = self.provider.create_stage( self.session, self.writer, stage=DatasourceStageInput( name="Retained register", source_name="retained_register", kind="upload", mode="static", shape="tabular", rows=({"id": 1},), governance=governance, ), ) datasource, first = self.provider.promote_stage( self.session, self.writer, stage_ref=first_stage.ref, ) second_stage = self.provider.create_stage( self.session, self.writer, stage=DatasourceStageInput( name="Retained register", source_name="retained_register", kind="upload", mode="static", shape="tabular", rows=({"id": 2},), target_datasource_ref=datasource.ref, ), ) _, second = self.provider.promote_stage( self.session, self.writer, stage_ref=second_stage.ref, ) as_of = max(first.created_at, second.created_at) + timedelta(days=2) plan = self.provider.preview_retention( self.session, principal("admin", scopes=("datasources:source:admin",)), as_of=as_of, ) by_ref = {item.ref: item for item in plan.candidates} self.assertTrue(by_ref[first.ref].eligible) self.assertFalse(by_ref[second.ref].eligible) self.assertIn("current_materialization", by_ref[second.ref].blockers) first_row = self.session.scalar( select(DatasourceMaterializationRecord).where( DatasourceMaterializationRecord.id == first.ref.removeprefix("materialization:") ) ) payload_id = first_row.payload_id disposed, evidence_hashes = self.provider.apply_retention( self.session, principal("admin", scopes=("datasources:source:admin",)), as_of=as_of, plan_hash=plan.plan_hash, target_refs=(first.ref,), ) self.assertEqual((first.ref,), disposed) self.assertEqual(1, len(evidence_hashes)) self.assertIsNone(self.session.get(DatasourcePayloadRecord, payload_id)) self.assertIsNotNone(first_row.disposed_at) self.assertEqual("disposed", first_row.state) with self.assertRaises(DatasourceUnavailableError): self.provider.read_datasource( self.session, self.writer, request=DatasourceReadRequest( datasource_ref=datasource.ref, materialization_ref=first.ref, ), ) held = self.provider.update_datasource_governance( self.session, self.writer, datasource_ref=datasource.ref, governance=DatasourceGovernance( retention_policy=governance.retention_policy, hold_refs=("hold:legal-1",), ), ) frozen = self.provider.freeze_datasource( self.session, self.writer, datasource_ref=held.ref, label="Legal evidence", ) held_plan = self.provider.preview_retention( self.session, principal("admin", scopes=("datasources:source:admin",)), as_of=frozen.created_at + timedelta(days=2), ) held_candidate = next(item for item in held_plan.candidates if item.ref == frozen.ref) self.assertFalse(held_candidate.eligible) self.assertIn("legal_hold", held_candidate.blockers) def test_retention_deletes_only_an_explicitly_selected_eligible_stage(self) -> None: stage = self.provider.create_stage( self.session, self.writer, stage=DatasourceStageInput( name="Transient import", source_name="transient_import", kind="upload", mode="static", shape="tabular", rows=({"id": 1},), governance=DatasourceGovernance( retention_policy={ "version": "stage-retention-v1", "enabled": True, "stage_days": 1, } ), ), ) admin = principal( "admin", scopes=("datasources:source:admin",), ) as_of = stage.created_at + timedelta(days=2) plan = self.provider.preview_retention( self.session, admin, as_of=as_of, ) candidate = next(item for item in plan.candidates if item.ref == stage.ref) self.assertTrue(candidate.eligible) self.assertEqual("delete_stage", candidate.disposition) disposed, evidence_hashes = self.provider.apply_retention( self.session, admin, as_of=as_of, plan_hash=plan.plan_hash, target_refs=(stage.ref,), ) self.assertEqual((stage.ref,), disposed) self.assertEqual(1, len(evidence_hashes)) self.assertIsNone( self.session.get( DatasourceStageRecord, stage.ref.removeprefix("stage:"), ) ) evidence = self.provider.list_lifecycle_evidence( self.session, self.writer, subject_ref=stage.ref, ) self.assertEqual("retention.stage_deleted", evidence[0].event_type) self.assertEqual(plan.plan_hash, evidence[0].details_["plan_hash"]) if __name__ == "__main__": unittest.main()