from __future__ import annotations import unittest from unittest.mock import patch 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, CAPABILITY_POLICY_DATASOURCE_VISIBILITY, DatasourceAccessError, DatasourceGovernance, DatasourceNotFoundError, DatasourceField, DatasourceOrigin, DatasourceOriginReadResult, DatasourceReadRequest, DatasourceVisibilityPolicyDecision, ) from govoplan_core.db.base import Base, utcnow from govoplan_datasources.backend.db.models import ( DatasourceGovernanceReferenceRecord, DatasourceMaterializationRecord, DatasourcePayloadRecord, DatasourceRecord, ) from govoplan_datasources.backend.service import ( CATALOGUE_READ_SCOPE, SqlDatasourceProvider, ) SCHEMA = [ {"name": "id", "data_type": "integer", "nullable": False}, {"name": "owner_id", "data_type": "string", "nullable": False}, {"name": "secret", "data_type": "string", "nullable": False}, {"name": "internal_note", "data_type": "string", "nullable": True}, ] ROWS = [ { "id": 1, "owner_id": "account-1", "secret": "protected-one", "internal_note": "hidden-one", }, { "id": 2, "owner_id": "account-2", "secret": "protected-two", "internal_note": "hidden-two", }, ] def principal( *, tenant_id: str = "tenant-1", account_id: str = "account-1", role_ids: tuple[str, ...] = ("reader",), auth_method: str = "session", service_account_id: str | None = None, ) -> ApiPrincipal: return ApiPrincipal( principal=PrincipalRef( account_id=account_id, membership_id=(None if auth_method == "service_account" else "member-1"), tenant_id=tenant_id, scopes=frozenset({CATALOGUE_READ_SCOPE}), role_ids=frozenset(role_ids), auth_method=auth_method, service_account_id=service_account_id, ), account=object(), user=object(), ) class _PolicyProvider: def __init__(self, decision: DatasourceVisibilityPolicyDecision) -> None: self.decision = decision self.requests = [] def decide_datasource_visibility(self, _session, *, request): self.requests.append(request) return self.decision class _Registry: def __init__(self, policy_provider: _PolicyProvider) -> None: self.policy_provider = policy_provider def has_capability(self, name: str) -> bool: return name == CAPABILITY_POLICY_DATASOURCE_VISIBILITY def capability(self, name: str) -> object: if not self.has_capability(name): raise KeyError(name) return self.policy_provider class _OriginProvider: origin = DatasourceOrigin( ref="secret:origin", source_name="live_cases", name="Live cases", kind="database", shape="tabular", supported_modes=("live",), provider="test", schema=tuple( DatasourceField( name=str(field["name"]), data_type=str(field["data_type"]), nullable=bool(field["nullable"]), ) for field in SCHEMA ), fingerprint="live-base-fingerprint", row_count=2, ) def list_origins(self, _session, _principal, *, query="", limit=100): del query return (self.origin,)[:limit] def get_origin(self, _session, _principal, *, origin_ref): return self.origin if origin_ref == self.origin.ref else None def read_origin(self, _session, _principal, *, request): rows = ROWS[request.offset : request.offset + request.limit] return DatasourceOriginReadResult( origin=self.origin, rows=tuple(rows), total_rows=2, truncated=request.offset + len(rows) < 2, elapsed_ms=1, ) class _LiveRegistry: def __init__(self) -> None: self.origin_provider = _OriginProvider() 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 DatasourceVisibilityTests(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__, DatasourceMaterializationRecord.__table__, ChangeSequenceEntry.__table__, ], ) self.Session = sessionmaker(bind=self.engine) self.session = self.Session() def tearDown(self) -> None: self.session.close() Base.metadata.drop_all( self.engine, tables=[ DatasourceMaterializationRecord.__table__, ChangeSequenceEntry.__table__, DatasourcePayloadRecord.__table__, DatasourceGovernanceReferenceRecord.__table__, DatasourceRecord.__table__, ], ) self.engine.dispose() def _datasource( self, *, policy: dict[str, object] | None = None, access_policy_ref: str | None = None, frozen: bool = False, snapshot_policy: dict[str, object] | None = None, ) -> DatasourceRecord: item = DatasourceRecord( tenant_id="tenant-1", source_name="governed_cases", name="Governed cases", kind="upload", mode="static", shape="tabular", status="active", provider="private.provider", provider_ref="secret:origin", schema_=SCHEMA, fingerprint="base-fingerprint", row_count=2, byte_count=250, provenance_={"secret_locator": "private://rows"}, metadata_={"credential_ref": "credential:secret"}, access_policy_ref=access_policy_ref, visibility_policy=policy or {}, ) self.session.add(item) self.session.flush() snapshot = DatasourceGovernance( publication_state="internal", visibility_policy=snapshot_policy or policy or {}, access_policy_ref=access_policy_ref, ) materialization = DatasourceMaterializationRecord( tenant_id="tenant-1", datasource_id=item.id, revision=1, state="published", schema_=SCHEMA, rows=ROWS, fingerprint="materialization-fingerprint", row_count=2, byte_count=250, frozen_at=utcnow() if frozen else None, frozen_label="Evidence" if frozen else None, provenance_={"protected": "source details"}, metadata_={"protected": "materialization details"}, governance_snapshot_=snapshot.to_dict(), ) self.session.add(materialization) self.session.flush() item.current_materialization_id = materialization.id self.session.flush() return item def test_role_acl_filters_discovery_and_denies_reads(self) -> None: item = self._datasource(policy={"source_acl": {"role_ids": ["reader"]}}) provider = SqlDatasourceProvider() self.assertEqual(1, len(provider.list_datasources(self.session, principal()))) denied = principal(role_ids=("other",)) self.assertEqual((), provider.list_datasources(self.session, denied)) self.assertIsNone( provider.get_datasource( self.session, denied, datasource_ref=f"datasource:{item.id}", ) ) with self.assertRaises(DatasourceAccessError): provider.read_datasource( self.session, denied, request=DatasourceReadRequest(datasource_ref=f"datasource:{item.id}"), ) other_tenant = principal(tenant_id="tenant-2") self.assertEqual((), provider.list_datasources(self.session, other_tenant)) with self.assertRaises(DatasourceNotFoundError): provider.read_datasource( self.session, other_tenant, request=DatasourceReadRequest(datasource_ref=f"datasource:{item.id}"), ) def test_rows_and_fields_are_filtered_into_an_opaque_permitted_view(self) -> None: policy = { "source_acl": {"role_ids": ["reader"]}, "fields": { "secret": { "classification": "restricted", "action": "redact", "allow": {"role_ids": ["privileged"]}, }, "internal_note": { "classification": "confidential", "action": "omit", "allow": {"role_ids": ["privileged"]}, }, }, "row_filters": [ {"field": "owner_id", "claim": "account_id", "operator": "equals"} ], } item = self._datasource(policy=policy) provider = SqlDatasourceProvider() with patch( "govoplan_datasources.backend.service.audit_event" ) as audit_event: result = provider.read_datasource( self.session, principal(), request=DatasourceReadRequest( datasource_ref=f"datasource:{item.id}" ), ) self.assertEqual( ({"id": 1, "owner_id": "account-1", "secret": None},), result.rows, ) self.assertEqual( ["id", "owner_id", "secret"], [field.name for field in result.datasource.schema], ) self.assertEqual("restricted", result.datasource.schema[-1].classification) self.assertNotEqual( "materialization-fingerprint", result.datasource.fingerprint ) self.assertIsNone(result.datasource.provider) self.assertIsNone(result.datasource.provider_ref) self.assertNotIn("credential_ref", result.datasource.metadata) self.assertEqual(1, result.total_rows) self.assertEqual( ["datasource.visibility_applied"], [diagnostic.code for diagnostic in result.diagnostics], ) audit_details = audit_event.call_args.kwargs["details"] self.assertEqual(2, audit_details["visibility_summary"]["field_rule_count"]) self.assertNotIn("protected-one", str(audit_details)) repeated = provider.read_datasource( self.session, principal(), request=DatasourceReadRequest( datasource_ref=f"datasource:{item.id}", expected_fingerprint=result.datasource.fingerprint, ), ) self.assertEqual(result.datasource.fingerprint, repeated.datasource.fingerprint) def test_materialization_acl_supports_service_principals(self) -> None: item = self._datasource( policy={ "source_acl": {"auth_methods": ["session", "service_account"]}, "materialization_acl": {"service_account_ids": ["service:reporting"]}, } ) provider = SqlDatasourceProvider() with self.assertRaises(DatasourceAccessError): provider.read_datasource( self.session, principal(), request=DatasourceReadRequest(datasource_ref=f"datasource:{item.id}"), ) result = provider.read_datasource( self.session, principal( account_id="service-account", auth_method="service_account", service_account_id="service:reporting", ), request=DatasourceReadRequest(datasource_ref=f"datasource:{item.id}"), ) self.assertEqual(2, result.total_rows) def test_frozen_state_keeps_its_restrictive_local_policy(self) -> None: snapshot_policy = {"source_acl": {"role_ids": ["evidence-reader"]}} item = self._datasource( policy={}, frozen=True, snapshot_policy=snapshot_policy, ) provider = SqlDatasourceProvider() with self.assertRaises(DatasourceAccessError): provider.read_datasource( self.session, principal(role_ids=("reader",)), request=DatasourceReadRequest( datasource_ref=f"datasource:{item.id}", consistency="frozen", ), ) result = provider.read_datasource( self.session, principal(role_ids=("evidence-reader",)), request=DatasourceReadRequest( datasource_ref=f"datasource:{item.id}", consistency="frozen", ), ) self.assertEqual(2, result.total_rows) def test_external_policy_tightens_local_behavior_and_missing_provider_denies( self, ) -> None: item = self._datasource(access_policy_ref="case-workers") with self.assertRaises(DatasourceAccessError): SqlDatasourceProvider().read_datasource( self.session, principal(), request=DatasourceReadRequest(datasource_ref=f"datasource:{item.id}"), ) external = _PolicyProvider( DatasourceVisibilityPolicyDecision( allowed=True, policies=({"source_acl": {"role_ids": ["case-worker"]}},), decision_ref="policy-decision:1", ) ) provider = SqlDatasourceProvider(registry=_Registry(external)) result = provider.read_datasource( self.session, principal(role_ids=("case-worker",)), request=DatasourceReadRequest(datasource_ref=f"datasource:{item.id}"), ) self.assertEqual(2, result.total_rows) self.assertEqual("case-workers", external.requests[-1].policy_ref) def test_denied_audit_contains_no_protected_values(self) -> None: item = self._datasource(policy={"source_acl": {"role_ids": ["authorized"]}}) provider = SqlDatasourceProvider() with patch("govoplan_datasources.backend.service.audit_event") as audit_event: with self.assertRaises(DatasourceAccessError): provider.read_datasource( self.session, principal(role_ids=("denied",)), request=DatasourceReadRequest( datasource_ref=f"datasource:{item.id}" ), ) details = audit_event.call_args.kwargs["details"] self.assertEqual("denied", details["outcome"]) serialized = str(details) self.assertNotIn("protected-one", serialized) self.assertNotIn("hidden-one", serialized) self.assertNotIn("credential:secret", serialized) def test_live_rows_are_filtered_before_origin_data_is_returned(self) -> None: item = self._datasource( policy={ "source_acl": {"role_ids": ["reader"]}, "row_filters": [ {"field": "owner_id", "claim": "account_id"} ], } ) item.mode = "live" self.session.flush() result = SqlDatasourceProvider(registry=_LiveRegistry()).read_datasource( self.session, principal(), request=DatasourceReadRequest( datasource_ref=f"datasource:{item.id}", consistency="live", ), ) self.assertEqual(1, result.total_rows) self.assertEqual("account-1", result.rows[0]["owner_id"]) self.assertNotEqual("live-base-fingerprint", result.datasource.fingerprint) self.assertIsNone(result.datasource.provider_ref) if __name__ == "__main__": unittest.main()