govoplan-datasources/tests/test_lifecycle_governance.py

from __future__ import annotations

from govoplan_core.core.tabular_sources import TabularCsvSource

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},),
                csv_source=TabularCsvSource(text="id\n1\n"),
                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
        self.assertIsNotNone(first_row.csv_source_)
        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.assertIsNone(first_row.csv_source_)
        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},),
                csv_source=TabularCsvSource(text="id\n1\n"),
                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()