from __future__ import annotations import json import uuid from typing import Any from sqlalchemy import select from app.approval.models import ApprovalRequest from app.approval.service import ApprovalService from app.audit.models import AuditLog from app.audit.service import AuditService from app.domains.data_access import constants as c from app.domains.data_access.models import ( BINDING_STATUS_ACTIVE, BINDING_STATUS_REVOKED, EFFECT_DESIRED_ABSENT, EFFECT_DESIRED_PRESENT, EFFECT_STATUS_APPLIED, ENTITLEMENT_STATUS_ACTIVE, ENTITLEMENT_VERSION_V1, OWNER_RESOLUTION_RESOLVED, OWNER_SOURCE_BOOTSTRAP, RECONCILE_STATUS_APPLIED, REQUEST_STATUS_APPROVED, RESOURCE_STATUS_ACTIVE, DataAccessBinding, DataAccessCompiledEffect, DataAccessEffectRun, DataAccessEntitlement, DataAccessRequest, DataAccessResource, ) from app.domains.data_access.reconciler import ( APPLY_ABSENT_OUTCOME_APPLIED, APPLY_PRESENT_OUTCOME_APPLIED, apply_absent_reconcile_compiled_effects, apply_present_reconcile_compiled_effects, ) from app.domains.data_access.service import DataAccessService from app.domains.email_group.models import EmailGroup, EmailGroupOwner from app.integrations.driver_base import ( DataAccessDatasetIamBinding, DataAccessDatasetIamPolicy, DataAccessDatasetIamPolicyPatch, DataAccessDatasetIamUpdateResult, DataAccessDriver, ) from app.integrations.models import Integration, IntegrationApprovalFlow class _FakeFeishuClient: def __init__(self) -> None: self.instances: list[dict[str, Any]] = [] async def get_approval_definition(self, approval_code: str) -> dict[str, Any]: return { "form": json.dumps([ {"id": "reason_widget", "name": "Business reason", "type": "textarea"}, {"id": "detail_widget", "name": "Detail URL", "type": "input"}, {"id": "datasets_widget", "name": "Datasets", "type": "textarea"}, {"id": "preview_widget", "name": "Compiled preview", "type": "textarea"}, ]), } async def resolve_user_id_by_email(self, email: str) -> str: return "feishu-user-dax-canary" async def create_instance( self, *, approval_code: str, user_id: str | None = None, open_id: str | None = None, form: str, instance_uuid: str | None = None, **kwargs: Any, ) -> dict[str, Any]: instance_code = f"dax-feishu-audit-{len(self.instances) + 1:03d}" self.instances.append({ "approval_code": approval_code, "user_id": user_id, "open_id": open_id, "form": form, "instance_uuid": instance_uuid, }) return { "instance_code": instance_code, "instance_status": "PENDING", } async def cancel_instance( self, *, instance_code: str, user_id: str, approval_code: str, ) -> dict[str, Any]: return {"ok": True} class _FakeApplyDriver(DataAccessDriver): provider = c.PROVIDER_BIGQUERY def __init__( self, *, policies: dict[tuple[str, str], DataAccessDatasetIamPolicy], apply_results: list[DataAccessDatasetIamUpdateResult], ) -> None: super().__init__( integration_id="fake-data-access-integration", config={}, ) self.policies = policies self.apply_results = list(apply_results) self.get_calls: list[tuple[str, str]] = [] self.set_calls: list[ tuple[str, str, str | None, DataAccessDatasetIamPolicyPatch] ] = [] self.forbidden_calls: list[str] = [] async def get_dataset_iam( self, project_id: str, dataset_id: str, ) -> DataAccessDatasetIamPolicy: self.get_calls.append((project_id, dataset_id)) return self.policies[(project_id, dataset_id)] async def set_dataset_iam( self, project_id: str, dataset_id: str, expected_etag: str | None, policy_patch: DataAccessDatasetIamPolicyPatch, ) -> DataAccessDatasetIamUpdateResult: self.set_calls.append((project_id, dataset_id, expected_etag, policy_patch)) return self.apply_results.pop(0) async def list_datasets(self, *args: Any, **kwargs: Any): self.forbidden_calls.append("list_datasets") raise AssertionError("DAX canary apply must not list datasets") async def list_all_datasets(self, *args: Any, **kwargs: Any): self.forbidden_calls.append("list_all_datasets") raise AssertionError("DAX canary apply must not list all datasets") async def list_resources(self, *args: Any, **kwargs: Any): self.forbidden_calls.append("list_resources") raise AssertionError("DAX canary apply must not list resources") async def list_all_resources(self, *args: Any, **kwargs: Any): self.forbidden_calls.append("list_all_resources") raise AssertionError("DAX canary apply must not list all resources") async def scan_resources(self, *args: Any, **kwargs: Any): self.forbidden_calls.append("scan_resources") raise AssertionError("DAX canary apply must not scan resources") async def get_project_iam(self, *args: Any, **kwargs: Any): self.forbidden_calls.append("get_project_iam") raise AssertionError("DAX canary apply must not read project IAM") async def set_project_iam(self, *args: Any, **kwargs: Any): self.forbidden_calls.append("set_project_iam") raise AssertionError("DAX canary apply must not write project IAM") async def get_project_policy(self, *args: Any, **kwargs: Any): self.forbidden_calls.append("get_project_policy") raise AssertionError("DAX canary apply must not read project policy") async def set_project_policy(self, *args: Any, **kwargs: Any): self.forbidden_calls.append("set_project_policy") raise AssertionError("DAX canary apply must not write project policy") async def grant_access(self, *args: Any, **kwargs: Any): self.forbidden_calls.append("grant_access") raise AssertionError("DAX canary apply must not use cloud grants") async def revoke_access(self, *args: Any, **kwargs: Any): self.forbidden_calls.append("revoke_access") raise AssertionError("DAX canary apply must not use cloud grants") async def _add_chat_flow(db_session) -> None: integration = Integration( kind="chat_approval", provider="feishu", name="Primary Feishu", slug=f"feishu-{uuid.uuid4().hex[:8]}", enabled=True, is_primary=True, config_public={"app_id": "unused-in-test"}, ) flow = IntegrationApprovalFlow( integration=integration, flow_key=c.APPROVAL_FLOW_DATA_ACCESS_GROUP_GRANT, label="Data Access Group Grant", external_code="approval-code-dax", payload_schema={ "request_summary": { "id": "widget-request-summary", "type": "textarea", }, "requested_access": { "id": "widget-requested-access", "type": "textarea", }, "dataset_list": { "id": "widget-dataset-list", "type": "textarea", }, "business_reason": { "id": "widget-business-reason", "type": "textarea", }, "detail_url": "widget-detail-url", }, active=True, ) db_session.add(flow) await db_session.flush() async def _add_data_access_integration(db_session) -> Integration: integration = Integration( kind=c.DOMAIN_DATA_ACCESS, provider=c.PROVIDER_BIGQUERY, name="Data Access BigQuery", slug=f"dax-bq-canary-{uuid.uuid4().hex[:8]}", enabled=True, is_primary=False, config_public={ "apply_mode_enabled": True, }, ) db_session.add(integration) await db_session.flush() return integration async def _add_entitlement(db_session) -> DataAccessEntitlement: entitlement = DataAccessEntitlement( domain=c.DOMAIN_DATA_ACCESS, provider=c.PROVIDER_BIGQUERY, resource_type=c.RESOURCE_TYPE_DATASET, principal_type=c.PRINCIPAL_TYPE_GOOGLE_GROUP, entitlement=c.ENTITLEMENT_READ, version=ENTITLEMENT_VERSION_V1, provider_member_prefix=c.BIGQUERY_GROUP_MEMBER_PREFIX, provider_role=c.BIGQUERY_DATA_VIEWER_ROLE, requestable=True, status=ENTITLEMENT_STATUS_ACTIVE, ) db_session.add(entitlement) await db_session.flush() return entitlement async def _add_resource( db_session, *, integration_id: uuid.UUID, owner_person_id: uuid.UUID, ) -> DataAccessResource: resource = DataAccessResource( integration_id=integration_id, provider=c.PROVIDER_BIGQUERY, resource_type=c.RESOURCE_TYPE_DATASET, project_id="helix-data-prod", dataset_id="finance", display_name="helix-data-prod.finance", requestable=True, status=RESOURCE_STATUS_ACTIVE, owner_person_id=str(owner_person_id), owner_resolution_status=OWNER_RESOLUTION_RESOLVED, owner_source=OWNER_SOURCE_BOOTSTRAP, source_version="bootstrap-2026-06-10", source_etag="etag-finance", ) db_session.add(resource) await db_session.flush() return resource async def _add_group( db_session, *, owner_person_id: uuid.UUID, ) -> EmailGroup: group = EmailGroup( email="canary-analytics@example.com", name="canary-analytics", status="active", google_group_id="google-group-canary-analytics", ) db_session.add(group) await db_session.flush() db_session.add( EmailGroupOwner( email_group_id=group.id, person_id=owner_person_id, ), ) await db_session.flush() return group async def _dispatch_approved_callback( db_session, approval_request: ApprovalRequest, ) -> None: await ApprovalService(db_session).handle_feishu_callback( {"event": {"instance_code": approval_request.external_id}}, prefetched_instance_detail={ "status": "APPROVED", "approval_code": "approval-code-dax", "instance_code": approval_request.external_id, }, ) async def _audit_by_action(db_session) -> dict[str, list[AuditLog]]: actions = { "data_access.request.submit", "data_access.request.callback_received", "data_access.binding.created", "data_access.compile.output", "data_access.request.approved", "data_access.reconcile.apply_present", "data_access.binding.revoke_requested", "data_access.reconcile.apply_absent", } audits = ( await db_session.scalars( select(AuditLog) .where(AuditLog.action.in_(actions)) .order_by(AuditLog.created_at, AuditLog.id), ) ).all() rows: dict[str, list[AuditLog]] = {action: [] for action in actions} for audit in audits: rows[audit.action].append(audit) return rows def _detail(audit: AuditLog) -> dict[str, Any]: assert audit.detail is not None return audit.detail def _assert_trace(audit: AuditLog, trace_id: str) -> None: assert audit.trace_id == AuditService._coerce_trace_id(trace_id) def _audit_with_phase(audits: list[AuditLog], phase: str) -> AuditLog: return next(audit for audit in audits if _detail(audit)["phase"] == phase) async def test_canary_request_apply_revoke_audit_trail_is_end_to_end( db_session, employee_user, ) -> None: requester, _ = employee_user await _add_chat_flow(db_session) integration = await _add_data_access_integration(db_session) await _add_entitlement(db_session) resource = await _add_resource( db_session, integration_id=integration.id, owner_person_id=requester.id, ) group = await _add_group(db_session, owner_person_id=requester.id) submission = await DataAccessService().submit_group_grant_request( db_session, requester=requester, resource_ids=[resource.id], principal=group.id, business_reason="Canary audit trail proof", detail_url="https://helix.example/data-access/canary-audit", idempotency_key=f"dax-037-{uuid.uuid4().hex}", trace_id="trace-dax-037-submit", feishu_client=_FakeFeishuClient(), ) assert submission.approval_request_id is not None request_row = await db_session.get(DataAccessRequest, submission.request_id) approval_request = await db_session.get( ApprovalRequest, submission.approval_request_id, ) assert request_row is not None assert approval_request is not None await _dispatch_approved_callback(db_session, approval_request) await db_session.refresh(request_row) assert request_row.status == REQUEST_STATUS_APPROVED binding = await db_session.scalar( select(DataAccessBinding).where( DataAccessBinding.source_request_id == request_row.id, ), ) assert binding is not None assert binding.status == BINDING_STATUS_ACTIVE present_effect = await db_session.scalar( select(DataAccessCompiledEffect).where( DataAccessCompiledEffect.binding_id == binding.id, DataAccessCompiledEffect.desired_state == EFFECT_DESIRED_PRESENT, ), ) assert present_effect is not None assert present_effect.provider_member == "group:canary-analytics@example.com" assert present_effect.provider_role == c.BIGQUERY_DATA_VIEWER_ROLE present_driver = _FakeApplyDriver( policies={ ("helix-data-prod", "finance"): DataAccessDatasetIamPolicy( project_id="helix-data-prod", dataset_id="finance", etag="provider-etag-present-1", bindings=(), raw_policy={"etag": "provider-etag-present-1"}, ), }, apply_results=[ DataAccessDatasetIamUpdateResult( ok=True, changed=True, etag="provider-etag-present-2", ), ], ) present_result = await apply_present_reconcile_compiled_effects( db_session, integration_id=integration.id, driver=present_driver, trace_id="trace-dax-037-apply-present", ) assert present_driver.forbidden_calls == [] assert present_driver.get_calls == [("helix-data-prod", "finance")] assert present_driver.set_calls == [( "helix-data-prod", "finance", "provider-etag-present-1", DataAccessDatasetIamPolicyPatch( add_bindings=( DataAccessDatasetIamBinding( role=c.BIGQUERY_DATA_VIEWER_ROLE, members=("group:canary-analytics@example.com",), ), ), remove_bindings=(), ), )] assert present_result.provider_write_count == 1 assert present_result.applied_count == 1 assert present_result.results[0].outcome == APPLY_PRESENT_OUTCOME_APPLIED revoke = await DataAccessService().revoke_data_access_binding( db_session, actor=requester, binding_id=binding.id, reason="canary_audit_revoke", trace_id="trace-dax-037-revoke", ) assert revoke.effect_count == 1 absent_effect = await db_session.scalar( select(DataAccessCompiledEffect).where( DataAccessCompiledEffect.binding_id == binding.id, DataAccessCompiledEffect.desired_state == EFFECT_DESIRED_ABSENT, ), ) assert absent_effect is not None assert absent_effect.provider_member == present_effect.provider_member assert absent_effect.provider_role == present_effect.provider_role absent_driver = _FakeApplyDriver( policies={ ("helix-data-prod", "finance"): DataAccessDatasetIamPolicy( project_id="helix-data-prod", dataset_id="finance", etag="provider-etag-absent-1", bindings=( DataAccessDatasetIamBinding( role=c.BIGQUERY_DATA_VIEWER_ROLE, members=( "group:canary-analytics@example.com", "group:unrelated-analytics@example.com", ), ), ), raw_policy={"etag": "provider-etag-absent-1"}, ), }, apply_results=[ DataAccessDatasetIamUpdateResult( ok=True, changed=True, etag="provider-etag-absent-2", ), ], ) absent_result = await apply_absent_reconcile_compiled_effects( db_session, integration_id=integration.id, driver=absent_driver, trace_id="trace-dax-037-apply-absent", ) assert absent_driver.forbidden_calls == [] assert absent_driver.get_calls == [("helix-data-prod", "finance")] assert absent_driver.set_calls == [( "helix-data-prod", "finance", "provider-etag-absent-1", DataAccessDatasetIamPolicyPatch( add_bindings=(), remove_bindings=( DataAccessDatasetIamBinding( role=c.BIGQUERY_DATA_VIEWER_ROLE, members=("group:canary-analytics@example.com",), ), ), ), )] assert absent_result.provider_write_count == 1 assert absent_result.removed_count == 1 assert absent_result.results[0].outcome == APPLY_ABSENT_OUTCOME_APPLIED await db_session.refresh(binding) await db_session.refresh(present_effect) await db_session.refresh(absent_effect) await db_session.refresh(request_row) assert binding.status == BINDING_STATUS_REVOKED assert binding.reconcile_status == RECONCILE_STATUS_APPLIED assert request_row.reconcile_status == RECONCILE_STATUS_APPLIED assert present_effect.status == EFFECT_STATUS_APPLIED assert absent_effect.status == EFFECT_STATUS_APPLIED runs = ( await db_session.scalars( select(DataAccessEffectRun) .where(DataAccessEffectRun.binding_id == binding.id) .order_by(DataAccessEffectRun.mode), ) ).all() assert {(run.mode, run.provider_outcome) for run in runs} == { ("apply_present", APPLY_PRESENT_OUTCOME_APPLIED), ("apply_absent", APPLY_ABSENT_OUTCOME_APPLIED), } assert { run.trace_id for run in runs } == { "trace-dax-037-apply-present", "trace-dax-037-apply-absent", } assert { run.provider_etag for run in runs } == {"provider-etag-present-2", "provider-etag-absent-2"} audits = await _audit_by_action(db_session) required_actions = { "data_access.request.submit", "data_access.request.callback_received", "data_access.binding.created", "data_access.compile.output", "data_access.request.approved", "data_access.reconcile.apply_present", "data_access.binding.revoke_requested", "data_access.reconcile.apply_absent", } assert {action for action, rows in audits.items() if rows} == required_actions assert all( audit.trace_id is not None for rows in audits.values() for audit in rows ) submit_audit = audits["data_access.request.submit"][0] _assert_trace(submit_audit, "trace-dax-037-submit") assert _detail(submit_audit)["dataset_count"] == 1 assert _detail(submit_audit)["feishu_instance_code"] == ( approval_request.external_id ) compile_audits = audits["data_access.compile.output"] submission_compile = _audit_with_phase(compile_audits, "submission_preview") _assert_trace(submission_compile, "trace-dax-037-submit") assert _detail(submission_compile)["pre_approval_preview_only"] is True assert _detail(submission_compile)["effects"][0]["desired_state"] == ( EFFECT_DESIRED_PRESENT ) callback_trace = f"data_access_callback:{approval_request.external_id}" callback_audit = audits["data_access.request.callback_received"][0] _assert_trace(callback_audit, callback_trace) assert callback_audit.resource_id == str(request_row.id) assert _detail(callback_audit)["request_id_raw"] == str(request_row.id) assert _detail(callback_audit)["approval_status"] == "approved" binding_audit = audits["data_access.binding.created"][0] _assert_trace(binding_audit, callback_trace) assert binding_audit.resource_id == str(binding.id) assert _detail(binding_audit)["request_id"] == str(request_row.id) assert _detail(binding_audit)["binding_status"] == BINDING_STATUS_ACTIVE callback_compile = _audit_with_phase(compile_audits, "approval_callback") _assert_trace(callback_compile, callback_trace) assert _detail(callback_compile)["effect_count"] == 1 assert _detail(callback_compile)["effects"][0]["effect_id"] == ( str(present_effect.id) ) approved_audit = audits["data_access.request.approved"][0] _assert_trace(approved_audit, callback_trace) assert _detail(approved_audit)["binding_count"] == 1 assert _detail(approved_audit)["effect_count"] == 1 present_audit = audits["data_access.reconcile.apply_present"][0] _assert_trace(present_audit, "trace-dax-037-apply-present") assert _detail(present_audit)["provider_mutation_attempted"] is True assert _detail(present_audit)["provider_write_count"] == 1 assert _detail(present_audit)["applied_count"] == 1 assert _detail(present_audit)["results"][0]["outcome"] == ( APPLY_PRESENT_OUTCOME_APPLIED ) assert _detail(present_audit)["results"][0]["provider_etag"] == ( "provider-etag-present-2" ) revoke_compile = _audit_with_phase(compile_audits, "revoke_request") _assert_trace(revoke_compile, "trace-dax-037-revoke") assert _detail(revoke_compile)["binding_id"] == str(binding.id) assert _detail(revoke_compile)["effects"][0]["desired_state"] == ( EFFECT_DESIRED_ABSENT ) assert _detail(revoke_compile)["effects"][0]["effect_id"] == str(absent_effect.id) revoke_audit = audits["data_access.binding.revoke_requested"][0] _assert_trace(revoke_audit, "trace-dax-037-revoke") assert _detail(revoke_audit)["reason"] == "canary_audit_revoke" assert _detail(revoke_audit)["effect_count"] == 1 absent_audit = audits["data_access.reconcile.apply_absent"][0] _assert_trace(absent_audit, "trace-dax-037-apply-absent") assert _detail(absent_audit)["provider_mutation_attempted"] is True assert _detail(absent_audit)["provider_write_count"] == 1 assert _detail(absent_audit)["removed_count"] == 1 assert _detail(absent_audit)["results"][0]["outcome"] == ( APPLY_ABSENT_OUTCOME_APPLIED ) assert _detail(absent_audit)["results"][0]["provider_etag"] == ( "provider-etag-absent-2" )