from __future__ import annotations import hashlib import json import re from collections.abc import Mapping, Sequence from dataclasses import dataclass, replace from typing import Any, cast from sqlalchemy import exists, func, or_, select, text from sqlalchemy.orm import Session from govoplan_core.audit.logging import audit_event from govoplan_core.auth import ApiPrincipal, has_scope from govoplan_core.core.datasources import ( DatasourceAccessError, DatasourceDescriptor, DatasourceError, DatasourceField, DatasourceGovernance, DatasourceMaterialization, DatasourceMode, DatasourceNotFoundError, DatasourceOrigin, DatasourceOriginReadRequest, DatasourcePublicationRequest, DatasourcePublicationResult, DatasourceReadRequest, DatasourceReadResult, DatasourceStage, DatasourceStageInput, DatasourceUnavailableError, DatasourceValidationError, datasource_origins, ) from govoplan_core.db.base import utcnow from govoplan_datasources.backend.db.models import ( DatasourceGovernanceReferenceRecord, DatasourceMaterializationRecord, DatasourcePayloadRecord, DatasourcePublicationRecord, DatasourceRecord, DatasourceStageRecord, ) from govoplan_datasources.backend.payloads import ( DatasourcePayloadBackend, PayloadBackendRegistry, create_database_rows_payload, payload_for_materialization, validate_payload_size, ) from govoplan_datasources.backend.quality import validate_stage from govoplan_datasources.backend.tabular import ( MAX_READ_ROWS, MAX_STAGE_ROWS, encoded_size, field_payload, fingerprint_rows, infer_schema, normalize_rows, ) CATALOGUE_READ_SCOPE = "datasources:catalogue:read" SOURCE_WRITE_SCOPE = "datasources:source:write" STAGE_WRITE_SCOPE = "datasources:stage:write" ADMIN_SCOPE = "datasources:source:admin" @dataclass(frozen=True, slots=True) class _PreparedPublication: producer_module: str producer_run_ref: str idempotency_key: str rows: tuple[dict[str, Any], ...] schema: tuple[DatasourceField, ...] fingerprint: str byte_count: int request_hash: str class SqlDatasourceProvider: def __init__( self, *, registry: object | None = None, payload_backends: Sequence[DatasourcePayloadBackend] = (), ) -> None: self._registry = registry self._payload_backends = PayloadBackendRegistry(payload_backends) def publish_rows( self, session: object, principal: object, *, request: DatasourcePublicationRequest, ) -> DatasourcePublicationResult: db, api_principal = _publication_context(session, principal) prepared = _prepare_publication(request) _lock_publication_identity( db, tenant_id=api_principal.tenant_id, producer_module=prepared.producer_module, idempotency_key=prepared.idempotency_key, ) existing = _existing_publication_result( db, tenant_id=api_principal.tenant_id, producer_module=prepared.producer_module, idempotency_key=prepared.idempotency_key, request_hash=prepared.request_hash, ) if existing is not None: return existing actor_id = _actor_id(api_principal) target, governance, baseline_schema = _publication_validation_context( db, tenant_id=api_principal.tenant_id, request=request, ) validation = validate_stage( rows=prepared.rows, schema=prepared.schema, quality_policy=governance.quality_policy, baseline_schema=baseline_schema, ) if validation["valid"] is not True: errors = validation.get("errors", []) error_codes = ", ".join( str(item.get("code") or "validation.error") for item in errors if isinstance(item, Mapping) ) suffix = f" ({error_codes})" if error_codes else "" raise DatasourceValidationError( "Published output failed governed quality or schema validation" f"{suffix}." ) datasource = _publication_target( db, tenant_id=api_principal.tenant_id, actor_id=actor_id, request=request, prepared=prepared, target=target, governance=governance, ) materialization = _append_materialization( db, datasource=datasource, rows=prepared.rows, schema=[field_payload(field) for field in prepared.schema], fingerprint=prepared.fingerprint, byte_count=prepared.byte_count, actor_id=actor_id, frozen=request.freeze, frozen_label=request.frozen_label, source_timestamp=request.source_timestamp, provenance={ **_publication_provenance(request, prepared), "publication_validation": validation, }, metadata=dict(request.metadata), set_current=request.set_current, ) publication = _create_publication_record( db, tenant_id=api_principal.tenant_id, actor_id=actor_id, datasource=datasource, materialization=materialization, request=request, prepared=prepared, ) audit_event( db, tenant_id=api_principal.tenant_id, user_id=getattr(api_principal.user, "id", None) or api_principal.account_id, api_key_id=api_principal.api_key_id, action="datasource.publication.published", object_type="datasource_publication", object_id=publication.id, details={ "producer_module": prepared.producer_module, "producer_run_ref": prepared.producer_run_ref, "datasource_ref": _datasource_ref(datasource.id), "materialization_ref": _materialization_ref(materialization.id), "fingerprint": prepared.fingerprint, "row_count": len(prepared.rows), "policy_hash": validation["policy_hash"], "schema_classification": validation["schema_change"][ "classification" ], }, ) return DatasourcePublicationResult( ref=_publication_ref(publication.id), status=publication.status, datasource=_datasource_dto(datasource), materialization=_materialization_dto(materialization), replayed=False, ) def list_datasources( self, session: object, principal: object, *, query: str = "", limit: int = 100, authority_mode: str | None = None, classification: str | None = None, publication_state: str | None = None, owner_ref: str | None = None, responsible_organization_ref: str | None = None, affected_ref: str | None = None, dependency_ref: str | None = None, ) -> Sequence[DatasourceDescriptor]: db, api_principal = _context(session, principal, CATALOGUE_READ_SCOPE) statement = ( select(DatasourceRecord) .where( DatasourceRecord.tenant_id == api_principal.tenant_id, DatasourceRecord.deleted_at.is_(None), ) .order_by(DatasourceRecord.name.asc(), DatasourceRecord.id.asc()) .limit(max(1, min(int(limit), 100))) ) cleaned_query = query.strip() if cleaned_query: pattern = f"%{_escape_like(cleaned_query)}%" statement = statement.where( or_( DatasourceRecord.name.ilike(pattern, escape="\\"), DatasourceRecord.source_name.ilike(pattern, escape="\\"), DatasourceRecord.description.ilike(pattern, escape="\\"), ) ) for column, value in ( (DatasourceRecord.authority_mode, authority_mode), (DatasourceRecord.classification, classification), (DatasourceRecord.publication_state, publication_state), (DatasourceRecord.owner_ref, owner_ref), ( DatasourceRecord.responsible_organization_ref, responsible_organization_ref, ), ): cleaned = str(value or "").strip() if cleaned: statement = statement.where(column == cleaned) for relation, value in ( ("affected", affected_ref), ("depends_on", dependency_ref), ): cleaned = str(value or "").strip() if cleaned: statement = statement.where( exists().where( DatasourceGovernanceReferenceRecord.datasource_id == DatasourceRecord.id, DatasourceGovernanceReferenceRecord.tenant_id == api_principal.tenant_id, DatasourceGovernanceReferenceRecord.relation == relation, DatasourceGovernanceReferenceRecord.reference == cleaned, ) ) return tuple(_datasource_dto(item) for item in db.scalars(statement)) def get_datasource( self, session: object, principal: object, *, datasource_ref: str, ) -> DatasourceDescriptor | None: db, api_principal = _context(session, principal, CATALOGUE_READ_SCOPE) item = _datasource_record( db, tenant_id=api_principal.tenant_id, datasource_ref=datasource_ref, ) return _datasource_dto(item) if item is not None else None def read_datasource( self, session: object, principal: object, *, request: DatasourceReadRequest, ) -> DatasourceReadResult: db, api_principal = _context(session, principal, CATALOGUE_READ_SCOPE) item = _required_datasource( db, tenant_id=api_principal.tenant_id, datasource_ref=request.datasource_ref, ) limit = max(1, min(int(request.limit), MAX_READ_ROWS)) offset = max(0, int(request.offset)) columns = tuple(dict.fromkeys(request.columns)) read_live = request.consistency == "live" or ( item.mode == "live" and not request.materialization_ref and request.consistency == "current" ) if read_live: if not item.provider_ref: raise DatasourceUnavailableError( "This datasource has no live origin." ) result = self._read_live( db, api_principal, item=item, request=request, limit=limit, offset=offset, columns=columns, ) if ( request.expected_fingerprint and request.expected_fingerprint != result.datasource.fingerprint ): raise DatasourceValidationError( "The live datasource fingerprint changed; refresh the consuming definition." ) return result return self._read_materialized( db, item=item, request=request, limit=limit, offset=offset, columns=columns, ) def _read_materialized( self, session: Session, *, item: DatasourceRecord, request: DatasourceReadRequest, limit: int, offset: int, columns: tuple[str, ...], ) -> DatasourceReadResult: materialization = _selected_materialization( session, item=item, request=request, ) if materialization is None: message = ( "No frozen state is available for this datasource." if request.consistency == "frozen" else "This datasource has no materialized state." ) raise DatasourceUnavailableError(message) if ( request.expected_fingerprint and request.expected_fingerprint != materialization.fingerprint ): raise DatasourceValidationError( "The datasource fingerprint changed; refresh the consuming definition." ) _validate_columns(materialization.schema_, columns) payload = payload_for_materialization(session, materialization) if payload is None: window = materialization.rows[offset : offset + limit] else: backend = self._payload_backends.require(payload.backend) backend.verify(session, payload) window = backend.read_rows( session, payload, offset=offset, limit=limit, ) rows = tuple(_select_columns(row, columns) for row in window) descriptor = _datasource_dto(item) return DatasourceReadResult( datasource=replace( descriptor, schema=_fields(materialization.schema_), schema_version=str(materialization.schema_version), fingerprint=materialization.fingerprint, row_count=materialization.row_count, byte_count=materialization.byte_count, ), rows=rows, total_rows=materialization.row_count, truncated=offset + len(rows) < materialization.row_count, materialization=_materialization_dto(materialization), ) def list_materializations( self, session: object, principal: object, *, datasource_ref: str, ) -> Sequence[DatasourceMaterialization]: db, api_principal = _context(session, principal, CATALOGUE_READ_SCOPE) item = _required_datasource( db, tenant_id=api_principal.tenant_id, datasource_ref=datasource_ref, ) statement = ( select(DatasourceMaterializationRecord) .where( DatasourceMaterializationRecord.tenant_id == api_principal.tenant_id, DatasourceMaterializationRecord.datasource_id == item.id, ) .order_by(DatasourceMaterializationRecord.revision.desc()) ) return tuple(_materialization_dto(row) for row in db.scalars(statement)) def list_stages( self, session: object, principal: object, *, limit: int = 100, ) -> Sequence[DatasourceStage]: db, api_principal = _context( session, principal, (CATALOGUE_READ_SCOPE, STAGE_WRITE_SCOPE), ) statement = ( select(DatasourceStageRecord) .where(DatasourceStageRecord.tenant_id == api_principal.tenant_id) .order_by(DatasourceStageRecord.created_at.desc()) .limit(max(1, min(int(limit), 100))) ) return tuple(_stage_dto(row) for row in db.scalars(statement)) def create_stage( self, session: object, principal: object, *, stage: DatasourceStageInput, ) -> DatasourceStage: db, api_principal = _context(session, principal, STAGE_WRITE_SCOPE) name = stage.name.strip() source_name = stage.source_name.strip() if not name: raise DatasourceValidationError("Stage name is required.") if not source_name: raise DatasourceValidationError("Datasource key is required.") if stage.shape != "tabular": raise DatasourceValidationError( "The first staging provider supports tabular data only." ) if stage.mode not in {"static", "cached"}: raise DatasourceValidationError( "Uploaded stages can be promoted as static or cached data." ) target = None if stage.target_datasource_ref: target = _required_datasource( db, tenant_id=api_principal.tenant_id, datasource_ref=stage.target_datasource_ref, ) if target.shape != stage.shape or target.mode != stage.mode: raise DatasourceValidationError( "A stage can only update a datasource with the same mode and shape." ) governance = ( stage.governance or (_datasource_governance(target) if target is not None else None) or _default_governance( mode=stage.mode, provider_ref=stage.provider_ref, ) ) rows = normalize_rows(stage.rows) schema = infer_schema(rows) fingerprint = fingerprint_rows(rows, schema) validation = validate_stage( rows=rows, schema=schema, quality_policy=governance.quality_policy, baseline_schema=_fields(target.schema_) if target is not None else None, ) item = DatasourceStageRecord( tenant_id=api_principal.tenant_id, target_datasource_id=target.id if target else None, name=name, source_name=source_name, description=_clean_optional(stage.description), kind=stage.kind, mode=stage.mode, shape=stage.shape, state="ready" if validation["valid"] else "invalid", provider=_clean_optional(stage.provider), provider_ref=_clean_optional(stage.provider_ref), schema_=[field_payload(field) for field in schema], rows=list(rows), fingerprint=fingerprint, row_count=len(rows), byte_count=encoded_size(rows), validation_=validation, provenance_=dict(stage.provenance), metadata_=dict(stage.metadata), governance_=governance.to_dict(), created_by=_actor_id(api_principal), ) db.add(item) db.flush() return _stage_dto(item) def promote_stage( self, session: object, principal: object, *, stage_ref: str, freeze: bool = False, frozen_label: str | None = None, ) -> tuple[DatasourceDescriptor, DatasourceMaterialization]: db, api_principal = _context(session, principal, STAGE_WRITE_SCOPE) stage = _required_stage( db, tenant_id=api_principal.tenant_id, stage_ref=stage_ref, ) if stage.validation_.get("valid") is not True: raise DatasourceValidationError( "The stage has blocking quality or schema diagnostics and cannot be promoted." ) if stage.state != "ready": raise DatasourceValidationError("Only ready stages can be promoted.") datasource = ( db.get(DatasourceRecord, stage.target_datasource_id) if stage.target_datasource_id else None ) if datasource is not None and ( datasource.tenant_id != api_principal.tenant_id or datasource.deleted_at is not None ): datasource = None if datasource is None: _ensure_source_name_available( db, tenant_id=api_principal.tenant_id, source_name=stage.source_name, ) datasource = DatasourceRecord( tenant_id=api_principal.tenant_id, source_name=stage.source_name, name=stage.name, description=stage.description, kind=stage.kind, mode=stage.mode, shape=stage.shape, status="active", provider=stage.provider or "datasources.stage", provider_ref=stage.provider_ref, schema_version=1, schema_=list(stage.schema_), fingerprint=stage.fingerprint, row_count=stage.row_count, byte_count=stage.byte_count, provenance_=dict(stage.provenance_), metadata_=dict(stage.metadata_), created_by=_actor_id(api_principal), updated_by=_actor_id(api_principal), ) _apply_datasource_governance( datasource, DatasourceGovernance.from_mapping(stage.governance_), ) db.add(datasource) db.flush() stage.target_datasource_id = datasource.id elif datasource.mode != stage.mode or datasource.shape != stage.shape: raise DatasourceValidationError( "A stage can only update a datasource with the same mode and shape." ) else: _apply_datasource_governance( datasource, DatasourceGovernance.from_mapping(stage.governance_), ) materialization = _append_materialization( db, datasource=datasource, rows=stage.rows, schema=stage.schema_, fingerprint=stage.fingerprint, byte_count=stage.byte_count, actor_id=_actor_id(api_principal), frozen=freeze, frozen_label=frozen_label, provenance={ **dict(stage.provenance_), "stage_ref": _stage_ref(stage.id), "stage_validation": dict(stage.validation_), }, metadata=dict(stage.metadata_), set_current=True, ) stage.state = "promoted" stage.promoted_at = utcnow() stage.promoted_materialization_id = materialization.id db.flush() return _datasource_dto(datasource), _materialization_dto(materialization) def register_origin( self, session: object, principal: object, *, origin_ref: str, name: str, source_name: str, mode: DatasourceMode, description: str | None = None, governance: DatasourceGovernance | None = None, ) -> DatasourceDescriptor: db, api_principal = _context(session, principal, SOURCE_WRITE_SCOPE) if mode not in {"live", "cached"}: raise DatasourceValidationError( "Connector origins can be registered as live or cached datasources." ) origin = self._required_origin(db, api_principal, origin_ref) if mode not in origin.supported_modes: raise DatasourceValidationError( f"The selected origin does not support {mode!r} datasources." ) cleaned_name = name.strip() cleaned_source_name = source_name.strip() if not cleaned_name or not cleaned_source_name: raise DatasourceValidationError( "Datasource name and key are required." ) _ensure_source_name_available( db, tenant_id=api_principal.tenant_id, source_name=cleaned_source_name, ) item = DatasourceRecord( tenant_id=api_principal.tenant_id, source_name=cleaned_source_name, name=cleaned_name, description=_clean_optional(description) or origin.description, kind=origin.kind, mode=mode, shape=origin.shape, status="active", provider=origin.provider, provider_ref=origin.ref, schema_version=_int_schema_version(origin.schema_version), schema_=[field_payload(field) for field in origin.schema], fingerprint=origin.fingerprint, row_count=origin.row_count, byte_count=origin.byte_count, provenance_={ "origin_ref": origin.ref, "origin_provider": origin.provider, "registered_at": utcnow().isoformat(), }, metadata_=_origin_metadata(origin), created_by=_actor_id(api_principal), updated_by=_actor_id(api_principal), ) _apply_datasource_governance( item, governance or _default_governance(mode=mode, provider_ref=origin.ref), ) db.add(item) db.flush() if mode == "cached": rows, refreshed_origin = self._read_origin_all( db, api_principal, origin_ref=origin_ref, ) self._materialize_origin( db, item=item, origin=refreshed_origin, rows=rows, actor_id=_actor_id(api_principal), set_current=True, ) return _datasource_dto(item) def update_datasource_governance( self, session: object, principal: object, *, datasource_ref: str, governance: DatasourceGovernance, ) -> DatasourceDescriptor: db, api_principal = _context(session, principal, SOURCE_WRITE_SCOPE) item = _required_datasource( db, tenant_id=api_principal.tenant_id, datasource_ref=datasource_ref, for_update=True, ) _apply_datasource_governance(item, governance) item.updated_by = _actor_id(api_principal) db.flush() return _datasource_dto(item) def refresh_datasource( self, session: object, principal: object, *, datasource_ref: str, ) -> tuple[DatasourceDescriptor, DatasourceMaterialization]: db, api_principal = _context(session, principal, SOURCE_WRITE_SCOPE) item = _required_datasource( db, tenant_id=api_principal.tenant_id, datasource_ref=datasource_ref, ) if item.mode != "cached" or not item.provider_ref: raise DatasourceValidationError( "Only cached connector-backed datasources can be refreshed." ) rows, origin = self._read_origin_all( db, api_principal, origin_ref=item.provider_ref, ) materialization = self._materialize_origin( db, item=item, origin=origin, rows=rows, actor_id=_actor_id(api_principal), set_current=True, ) return _datasource_dto(item), _materialization_dto(materialization) def freeze_datasource( self, session: object, principal: object, *, datasource_ref: str, label: str | None = None, ) -> DatasourceMaterialization: db, api_principal = _context(session, principal, SOURCE_WRITE_SCOPE) item = _required_datasource( db, tenant_id=api_principal.tenant_id, datasource_ref=datasource_ref, ) if item.mode == "live": if not item.provider_ref: raise DatasourceUnavailableError( "The live datasource no longer references an origin." ) rows, origin = self._read_origin_all( db, api_principal, origin_ref=item.provider_ref, ) materialization = self._materialize_origin( db, item=item, origin=origin, rows=rows, actor_id=_actor_id(api_principal), frozen=True, frozen_label=label, set_current=False, ) return _materialization_dto(materialization) current = _current_materialization(db, item) if current is None: raise DatasourceUnavailableError( "The datasource has no current state to freeze." ) current_payload = payload_for_materialization(db, current) materialization = _append_materialization( db, datasource=item, rows=current.rows if current_payload is None else (), schema=current.schema_, fingerprint=current.fingerprint, byte_count=current.byte_count, actor_id=_actor_id(api_principal), frozen=True, frozen_label=label, source_timestamp=current.source_timestamp, provenance={ **dict(current.provenance_), "frozen_from": _materialization_ref(current.id), }, metadata=dict(current.metadata_), set_current=False, reusable_payload=current_payload, ) return _materialization_dto(materialization) def retire_datasource( self, session: object, principal: object, *, datasource_ref: str, ) -> DatasourceDescriptor: db, api_principal = _context(session, principal, SOURCE_WRITE_SCOPE) item = _required_datasource( db, tenant_id=api_principal.tenant_id, datasource_ref=datasource_ref, ) item.status = "retired" item.deleted_at = utcnow() item.updated_by = _actor_id(api_principal) db.flush() return _datasource_dto(item) def _read_live( self, session: Session, principal: ApiPrincipal, *, item: DatasourceRecord, request: DatasourceReadRequest, limit: int, offset: int, columns: tuple[str, ...], ) -> DatasourceReadResult: if not item.provider_ref: raise DatasourceUnavailableError( "The live datasource no longer references an origin." ) provider = datasource_origins(self._registry) if provider is None: raise DatasourceUnavailableError( "No connector origin provider is currently available." ) try: result = provider.read_origin( session, principal, request=DatasourceOriginReadRequest( origin_ref=item.provider_ref, limit=limit, offset=offset, columns=columns, expected_fingerprint=request.expected_fingerprint, max_bytes=request.max_bytes, timeout_ms=request.timeout_ms, ), ) except DatasourceError: raise except Exception as exc: raise DatasourceUnavailableError( f"The datasource origin could not be read: {exc}" ) from exc descriptor = replace( _datasource_dto(item), schema=result.origin.schema, schema_version=result.origin.schema_version, fingerprint=result.origin.fingerprint, row_count=result.origin.row_count, byte_count=result.origin.byte_count, updated_at=result.origin.updated_at, metadata=_origin_metadata(result.origin, current=item.metadata_), ) return DatasourceReadResult( datasource=descriptor, rows=result.rows, total_rows=result.total_rows, truncated=result.truncated, returned_bytes=result.returned_bytes, elapsed_ms=result.elapsed_ms, effective_row_limit=result.effective_row_limit, effective_byte_limit=result.effective_byte_limit, effective_timeout_ms=result.effective_timeout_ms, diagnostics=result.diagnostics, ) def _required_origin( self, session: Session, principal: ApiPrincipal, origin_ref: str, ) -> DatasourceOrigin: provider = datasource_origins(self._registry) if provider is None: raise DatasourceUnavailableError( "No connector origin provider is currently available." ) try: origin = provider.get_origin( session, principal, origin_ref=origin_ref, ) except DatasourceError: raise except Exception as exc: raise DatasourceUnavailableError( f"The datasource origin could not be inspected: {exc}" ) from exc if origin is None: raise DatasourceNotFoundError("Datasource origin not found.") return origin def _read_origin_all( self, session: Session, principal: ApiPrincipal, *, origin_ref: str, ) -> tuple[tuple[Mapping[str, object], ...], DatasourceOrigin]: provider = datasource_origins(self._registry) if provider is None: raise DatasourceUnavailableError( "No connector origin provider is currently available." ) rows: list[Mapping[str, object]] = [] origin: DatasourceOrigin | None = None offset = 0 while True: try: result = provider.read_origin( session, principal, request=DatasourceOriginReadRequest( origin_ref=origin_ref, limit=min(MAX_READ_ROWS, MAX_STAGE_ROWS - offset), offset=offset, ), ) except DatasourceError: raise except Exception as exc: raise DatasourceUnavailableError( f"The datasource origin could not be materialized: {exc}" ) from exc origin = result.origin rows.extend(result.rows) offset += len(result.rows) if not result.truncated: break if offset >= MAX_STAGE_ROWS or not result.rows: raise DatasourceValidationError( f"Cached and frozen sources are limited to {MAX_STAGE_ROWS:,} rows." ) if origin is None: raise DatasourceUnavailableError("The datasource origin returned no metadata.") return tuple(rows), origin def _materialize_origin( self, session: Session, *, item: DatasourceRecord, origin: DatasourceOrigin, rows: Sequence[Mapping[str, object]], actor_id: str | None, frozen: bool = False, frozen_label: str | None = None, set_current: bool, ) -> DatasourceMaterializationRecord: item.metadata_ = _origin_metadata(origin, current=item.metadata_) normalized = normalize_rows(rows) schema = infer_schema(normalized) or origin.schema fingerprint = fingerprint_rows(normalized, schema) return _append_materialization( session, datasource=item, rows=normalized, schema=[field_payload(field) for field in schema], fingerprint=fingerprint, byte_count=encoded_size(normalized), actor_id=actor_id, frozen=frozen, frozen_label=frozen_label, source_timestamp=origin.updated_at, provenance={ **dict(item.provenance_), "origin_ref": origin.ref, "origin_fingerprint": origin.fingerprint, "materialized_at": utcnow().isoformat(), }, metadata={ **dict(item.metadata_), "origin_schema_version": origin.schema_version, }, set_current=set_current, ) def _origin_metadata( origin: DatasourceOrigin, *, current: Mapping[str, object] | None = None, ) -> dict[str, object]: return { **dict(current or {}), **dict(origin.metadata), "source_contract": { "source_mode": origin.source_mode, "pushdown": { "projections": origin.pushdown.projections, "pagination": origin.pushdown.pagination, "filters": list(origin.pushdown.filters), "aggregations": list(origin.pushdown.aggregations), "sorting": list(origin.pushdown.sorting), }, "health": { "status": origin.health.status, "code": origin.health.code, "summary": origin.health.summary, "checked_at": ( origin.health.checked_at.isoformat() if origin.health.checked_at else None ), "details": dict(origin.health.details), }, }, } def _append_materialization( session: Session, *, datasource: DatasourceRecord, rows: Sequence[Mapping[str, object]], schema: Sequence[Mapping[str, object]], fingerprint: str, byte_count: int, actor_id: str | None, frozen: bool = False, frozen_label: str | None = None, source_timestamp=None, provenance: Mapping[str, object] | None = None, metadata: Mapping[str, object] | None = None, set_current: bool, reusable_payload: DatasourcePayloadRecord | None = None, ) -> DatasourceMaterializationRecord: datasource = _lock_datasource_for_materialization(session, datasource) revision = _allocate_materialization_revision(session, datasource) payload = reusable_payload or create_database_rows_payload( session, tenant_id=datasource.tenant_id, rows=rows, actor_id=actor_id, metadata={ "datasource_id": datasource.id, "fingerprint": fingerprint, }, ) if payload.tenant_id != datasource.tenant_id or payload.state != "ready": raise DatasourceValidationError( "Only a ready payload from the same tenant can be materialized." ) validate_payload_size(payload, expected_byte_count=byte_count) schema_payload, schema_version = _materialization_schema( datasource, schema, ) materialization = DatasourceMaterializationRecord( tenant_id=datasource.tenant_id, datasource_id=datasource.id, revision=revision, state="published", schema_version=max(1, int(schema_version or 1)), schema_=schema_payload, payload_id=payload.id, payload_checksum=payload.checksum, rows=[], fingerprint=fingerprint, row_count=payload.row_count, byte_count=payload.byte_count, frozen_at=utcnow() if frozen else None, frozen_label=_clean_optional(frozen_label), source_timestamp=source_timestamp, provenance_=dict(provenance or {}), metadata_=dict(metadata or {}), governance_snapshot_=_datasource_governance(datasource).to_dict(), created_by=actor_id, ) session.add(materialization) session.flush() if set_current: _apply_current_materialization( datasource, materialization=materialization, schema=schema_payload, schema_version=schema_version, actor_id=actor_id, ) session.flush() return materialization def _lock_datasource_for_materialization( session: Session, datasource: DatasourceRecord, ) -> DatasourceRecord: locked = session.scalar( select(DatasourceRecord) .where( DatasourceRecord.id == datasource.id, DatasourceRecord.tenant_id == datasource.tenant_id, DatasourceRecord.deleted_at.is_(None), ) .with_for_update() .execution_options(populate_existing=True) ) if locked is None: raise DatasourceUnavailableError( "The datasource is no longer available for materialization." ) return locked def _allocate_materialization_revision( session: Session, datasource: DatasourceRecord, ) -> int: """Allocate under the datasource row lock held by the caller.""" return int( session.scalar( select(func.max(DatasourceMaterializationRecord.revision)).where( DatasourceMaterializationRecord.datasource_id == datasource.id ) ) or 0 ) + 1 def _materialization_schema( datasource: DatasourceRecord, schema: Sequence[Mapping[str, object]], ) -> tuple[list[dict[str, object]], int]: schema_payload = [dict(field) for field in schema] schema_changed = datasource.schema_ != schema_payload schema_version = ( datasource.schema_version + 1 if schema_changed else datasource.schema_version ) return schema_payload, int(schema_version or 1) def _apply_current_materialization( datasource: DatasourceRecord, *, materialization: DatasourceMaterializationRecord, schema: list[dict[str, object]], schema_version: int, actor_id: str | None, ) -> None: datasource.current_materialization_id = materialization.id datasource.schema_ = schema datasource.schema_version = schema_version datasource.fingerprint = materialization.fingerprint datasource.row_count = materialization.row_count datasource.byte_count = materialization.byte_count datasource.updated_by = actor_id def _selected_materialization( session: Session, *, item: DatasourceRecord, request: DatasourceReadRequest, ) -> DatasourceMaterializationRecord | None: if request.materialization_ref: return _materialization_record( session, tenant_id=item.tenant_id, datasource_id=item.id, materialization_ref=request.materialization_ref, ) if request.consistency == "frozen": return session.scalar( select(DatasourceMaterializationRecord) .where( DatasourceMaterializationRecord.datasource_id == item.id, DatasourceMaterializationRecord.tenant_id == item.tenant_id, DatasourceMaterializationRecord.frozen_at.is_not(None), ) .order_by(DatasourceMaterializationRecord.revision.desc()) .limit(1) ) return _current_materialization(session, item) def _current_materialization( session: Session, item: DatasourceRecord, ) -> DatasourceMaterializationRecord | None: if not item.current_materialization_id: return None return session.scalar( select(DatasourceMaterializationRecord).where( DatasourceMaterializationRecord.id == item.current_materialization_id, DatasourceMaterializationRecord.datasource_id == item.id, DatasourceMaterializationRecord.tenant_id == item.tenant_id, ) ) def _datasource_record( session: Session, *, tenant_id: str, datasource_ref: str, for_update: bool = False, ) -> DatasourceRecord | None: datasource_id = _strip_ref(datasource_ref, "datasource:") if datasource_id is None: return None statement = select(DatasourceRecord).where( DatasourceRecord.id == datasource_id, DatasourceRecord.tenant_id == tenant_id, DatasourceRecord.deleted_at.is_(None), ) if for_update: statement = statement.with_for_update() return session.scalar(statement) def _required_datasource( session: Session, *, tenant_id: str, datasource_ref: str, for_update: bool = False, ) -> DatasourceRecord: item = _datasource_record( session, tenant_id=tenant_id, datasource_ref=datasource_ref, for_update=for_update, ) if item is None: raise DatasourceNotFoundError("Datasource not found.") return item def _materialization_record( session: Session, *, tenant_id: str, datasource_id: str, materialization_ref: str, ) -> DatasourceMaterializationRecord | None: materialization_id = _strip_ref(materialization_ref, "materialization:") if materialization_id is None: return None return session.scalar( select(DatasourceMaterializationRecord).where( DatasourceMaterializationRecord.id == materialization_id, DatasourceMaterializationRecord.tenant_id == tenant_id, DatasourceMaterializationRecord.datasource_id == datasource_id, ) ) def _required_stage( session: Session, *, tenant_id: str, stage_ref: str, ) -> DatasourceStageRecord: stage_id = _strip_ref(stage_ref, "stage:") if stage_id is None: raise DatasourceNotFoundError("Datasource stage not found.") item = session.scalar( select(DatasourceStageRecord).where( DatasourceStageRecord.id == stage_id, DatasourceStageRecord.tenant_id == tenant_id, ) ) if item is None: raise DatasourceNotFoundError("Datasource stage not found.") return item def _ensure_source_name_available( session: Session, *, tenant_id: str, source_name: str, ) -> None: existing = session.scalar( select(DatasourceRecord.id).where( DatasourceRecord.tenant_id == tenant_id, DatasourceRecord.source_name == source_name, ) ) if existing is not None: raise DatasourceValidationError( f"A datasource with key {source_name!r} already exists." ) def _default_governance( *, mode: str, provider_ref: str | None, ) -> DatasourceGovernance: if mode == "live": authority_mode = "external_authoritative" elif mode == "cached" and provider_ref: authority_mode = "external_mirror" else: authority_mode = "native_authoritative" return DatasourceGovernance( authoritative_source_ref=provider_ref, authority_mode=cast(Any, authority_mode), purposes=("governed_data_processing",), publication_state="internal", ) def _datasource_governance(item: DatasourceRecord) -> DatasourceGovernance: return DatasourceGovernance.from_mapping( { "owner_ref": item.owner_ref, "steward_ref": item.steward_ref, "responsible_organization_ref": item.responsible_organization_ref, "responsible_function_ref": item.responsible_function_ref, "authoritative_source_ref": item.authoritative_source_ref, "authority_mode": item.authority_mode, "legal_basis_refs": item.legal_basis_refs, "purposes": item.purposes, "semantic_definition": item.semantic_definition, "schema_owner_ref": item.schema_owner_ref, "official_keys": item.official_keys, "classification": item.classification, "privacy_profile_ref": item.privacy_profile_ref, "retention_policy_ref": item.retention_policy_ref, "hold_refs": item.hold_refs, "publication_state": item.publication_state, "transfer_agreement_ref": item.transfer_agreement_ref, "freshness_policy": item.freshness_policy, "quality_policy": item.quality_policy, "known_limits": item.known_limits, "correction_procedure_ref": item.correction_procedure_ref, "affected_refs": item.affected_refs, "dependency_refs": item.dependency_refs, } ) def _apply_datasource_governance( item: DatasourceRecord, governance: DatasourceGovernance, ) -> None: item.owner_ref = governance.owner_ref item.steward_ref = governance.steward_ref item.responsible_organization_ref = governance.responsible_organization_ref item.responsible_function_ref = governance.responsible_function_ref item.authoritative_source_ref = governance.authoritative_source_ref item.authority_mode = governance.authority_mode item.legal_basis_refs = list(governance.legal_basis_refs) item.purposes = list(governance.purposes) item.semantic_definition = governance.semantic_definition item.schema_owner_ref = governance.schema_owner_ref item.official_keys = list(governance.official_keys) item.classification = governance.classification item.privacy_profile_ref = governance.privacy_profile_ref item.retention_policy_ref = governance.retention_policy_ref item.hold_refs = list(governance.hold_refs) item.publication_state = governance.publication_state item.transfer_agreement_ref = governance.transfer_agreement_ref item.freshness_policy = dict(governance.freshness_policy) item.quality_policy = dict(governance.quality_policy) item.known_limits = list(governance.known_limits) item.correction_procedure_ref = governance.correction_procedure_ref item.affected_refs = list(governance.affected_refs) item.dependency_refs = list(governance.dependency_refs) item.governance_references = [ DatasourceGovernanceReferenceRecord( tenant_id=item.tenant_id, relation=relation, reference=reference, ) for relation, references in ( ("affected", governance.affected_refs), ("depends_on", governance.dependency_refs), ) for reference in references ] def _datasource_dto(item: DatasourceRecord) -> DatasourceDescriptor: capabilities = ["read", "preview", "freeze"] if item.mode == "cached" and item.provider_ref: capabilities.append("refresh") return DatasourceDescriptor( ref=_datasource_ref(item.id), source_name=item.source_name, name=item.name, description=item.description, kind=cast(Any, item.kind), mode=cast(Any, item.mode), shape=cast(Any, item.shape), status=item.status, provider=item.provider, provider_ref=item.provider_ref, schema=_fields(item.schema_), schema_version=str(item.schema_version), fingerprint=item.fingerprint, current_materialization_ref=( _materialization_ref(item.current_materialization_id) if item.current_materialization_id else None ), row_count=item.row_count, byte_count=item.byte_count, updated_at=item.updated_at, capabilities=tuple(capabilities), provenance=dict(item.provenance_), metadata=dict(item.metadata_), governance=_datasource_governance(item), ) def _materialization_dto( item: DatasourceMaterializationRecord, ) -> DatasourceMaterialization: return DatasourceMaterialization( ref=_materialization_ref(item.id), datasource_ref=_datasource_ref(item.datasource_id), revision=item.revision, state=item.state, fingerprint=item.fingerprint, schema=_fields(item.schema_), row_count=item.row_count, byte_count=item.byte_count, frozen_at=item.frozen_at, frozen_label=item.frozen_label, source_timestamp=item.source_timestamp, created_at=item.created_at, provenance=dict(item.provenance_), metadata=dict(item.metadata_), governance=DatasourceGovernance.from_mapping(item.governance_snapshot_), ) def _stage_dto(item: DatasourceStageRecord) -> DatasourceStage: return DatasourceStage( ref=_stage_ref(item.id), name=item.name, source_name=item.source_name, kind=cast(Any, item.kind), mode=cast(Any, item.mode), shape=cast(Any, item.shape), state=item.state, target_datasource_ref=( _datasource_ref(item.target_datasource_id) if item.target_datasource_id else None ), fingerprint=item.fingerprint, schema=_fields(item.schema_), row_count=item.row_count, byte_count=item.byte_count, validation=dict(item.validation_), created_at=item.created_at, promoted_at=item.promoted_at, promoted_materialization_ref=( _materialization_ref(item.promoted_materialization_id) if item.promoted_materialization_id else None ), provenance=dict(item.provenance_), metadata=dict(item.metadata_), governance=DatasourceGovernance.from_mapping(item.governance_), ) def _fields(payload: Sequence[Mapping[str, object]]) -> tuple[DatasourceField, ...]: return tuple( DatasourceField( name=str(item.get("name", "")), data_type=str(item.get("data_type", "unknown")), nullable=bool(item.get("nullable", True)), ) for item in payload ) def _context( session: object, principal: object, required_scope: str | Sequence[str], ) -> tuple[Session, ApiPrincipal]: if not isinstance(session, Session): raise TypeError("Datasource providers require a SQLAlchemy session.") if not isinstance(principal, ApiPrincipal): raise DatasourceAccessError("A tenant API principal is required.") required_scopes = ( (required_scope,) if isinstance(required_scope, str) else tuple(required_scope) ) if not ( any(has_scope(principal, scope) for scope in required_scopes) or has_scope(principal, ADMIN_SCOPE) ): raise DatasourceAccessError( f"Missing one of the required scopes: {', '.join(required_scopes)}" ) return session, principal def _publication_context( session: object, principal: object, ) -> tuple[Session, ApiPrincipal]: return _context(session, principal, SOURCE_WRITE_SCOPE) def _validate_columns( schema: Sequence[Mapping[str, object]], columns: tuple[str, ...], ) -> None: if not columns: return known = {str(field.get("name", "")) for field in schema} unknown = [column for column in columns if column not in known] if unknown: raise DatasourceValidationError( f"Unknown datasource columns: {', '.join(unknown)}" ) def _select_columns( row: Mapping[str, object], columns: tuple[str, ...], ) -> Mapping[str, object]: if not columns: return dict(row) return {column: row.get(column) for column in columns} def _strip_ref(value: str, prefix: str) -> str | None: cleaned = str(value or "").strip() if not cleaned: return None if cleaned.startswith(prefix): return cleaned[len(prefix) :] return cleaned if ":" not in cleaned else None def _datasource_ref(item_id: str) -> str: return f"datasource:{item_id}" def _materialization_ref(item_id: str) -> str: return f"materialization:{item_id}" def _stage_ref(item_id: str) -> str: return f"stage:{item_id}" def _publication_ref(item_id: str) -> str: return f"publication:{item_id}" def _actor_id(principal: ApiPrincipal) -> str | None: return principal.account_id or principal.membership_id or principal.identity_id def _clean_optional(value: object | None) -> str | None: cleaned = str(value or "").strip() return cleaned or None def _int_schema_version(value: str) -> int: try: return max(1, int(value)) except (TypeError, ValueError): return 1 def _escape_like(value: str) -> str: return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") def _valid_source_name(value: str) -> bool: return bool(re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]{0,119}", value)) def _prepare_publication( request: DatasourcePublicationRequest, ) -> _PreparedPublication: producer_module = request.producer_module.strip() producer_run_ref = request.producer_run_ref.strip() idempotency_key = request.idempotency_key.strip() _validate_publication_identity( producer_module=producer_module, producer_run_ref=producer_run_ref, idempotency_key=idempotency_key, ) normalized = normalize_rows(request.rows) schema = infer_schema(normalized) fingerprint = fingerprint_rows(normalized, schema) return _PreparedPublication( producer_module=producer_module, producer_run_ref=producer_run_ref, idempotency_key=idempotency_key, rows=normalized, schema=schema, fingerprint=fingerprint, byte_count=encoded_size(normalized), request_hash=_publication_request_hash( request, normalized=normalized, fingerprint=fingerprint, ), ) def _validate_publication_identity( *, producer_module: str, producer_run_ref: str, idempotency_key: str, ) -> None: values = ( (producer_module, 100, "A producer module"), (producer_run_ref, 500, "A producer run reference"), (idempotency_key, 255, "An idempotency key"), ) for value, maximum, label in values: if not value or len(value) > maximum: raise DatasourceValidationError( f"{label} of at most {maximum} characters is required." ) def _publication_lock_key( *, tenant_id: str, producer_module: str, idempotency_key: str, ) -> int: payload = ( f"govoplan.datasources.publication\0{tenant_id}\0" f"{producer_module}\0{idempotency_key}" ).encode("utf-8") return int.from_bytes( hashlib.sha256(payload).digest()[:8], "big", signed=True, ) def _lock_publication_identity( session: Session, *, tenant_id: str, producer_module: str, idempotency_key: str, ) -> None: if session.get_bind().dialect.name != "postgresql": return session.execute( text("SELECT pg_advisory_xact_lock(:key)"), { "key": _publication_lock_key( tenant_id=tenant_id, producer_module=producer_module, idempotency_key=idempotency_key, ) }, ) def _existing_publication_result( session: Session, *, tenant_id: str, producer_module: str, idempotency_key: str, request_hash: str, ) -> DatasourcePublicationResult | None: publication = session.scalar( select(DatasourcePublicationRecord).where( DatasourcePublicationRecord.tenant_id == tenant_id, DatasourcePublicationRecord.producer_module == producer_module, DatasourcePublicationRecord.idempotency_key == idempotency_key, ) ) if publication is None: return None if publication.request_hash != request_hash: raise DatasourceValidationError( "The publication idempotency key was already used with different output." ) datasource = session.get(DatasourceRecord, publication.datasource_id) materialization = session.get( DatasourceMaterializationRecord, publication.materialization_id, ) if ( datasource is None or materialization is None or datasource.tenant_id != tenant_id or materialization.tenant_id != tenant_id ): raise DatasourceUnavailableError( "The prior publication result is no longer available." ) return DatasourcePublicationResult( ref=_publication_ref(publication.id), status=publication.status, datasource=_datasource_dto(datasource), materialization=_materialization_dto(materialization), replayed=True, ) def _publication_target( session: Session, *, tenant_id: str, actor_id: str | None, request: DatasourcePublicationRequest, prepared: _PreparedPublication, target: DatasourceRecord | None, governance: DatasourceGovernance, ) -> DatasourceRecord: if target is not None: datasource = target if request.governance is not None: _apply_datasource_governance(datasource, governance) datasource.updated_by = actor_id return datasource name = str(request.name or "").strip() source_name = str(request.source_name or "").strip() if not name: raise DatasourceValidationError( "A datasource name is required for a new publication target." ) if not _valid_source_name(source_name): raise DatasourceValidationError( "Datasource keys must start with a letter or underscore and contain " "only letters, numbers, and underscores." ) _ensure_source_name_available( session, tenant_id=tenant_id, source_name=source_name, ) datasource = DatasourceRecord( tenant_id=tenant_id, source_name=source_name, name=name, description=_clean_optional(request.description), kind="custom", mode="static", shape="tabular", status="active", provider=prepared.producer_module, provider_ref=prepared.producer_run_ref, schema_version=1, schema_=[field_payload(field) for field in prepared.schema], fingerprint=prepared.fingerprint, row_count=len(prepared.rows), byte_count=prepared.byte_count, provenance_={ **dict(request.provenance), "producer_module": prepared.producer_module, "producer_run_ref": prepared.producer_run_ref, }, metadata_=dict(request.metadata), created_by=actor_id, updated_by=actor_id, ) _apply_datasource_governance( datasource, governance, ) session.add(datasource) session.flush() return datasource def _publication_validation_context( session: Session, *, tenant_id: str, request: DatasourcePublicationRequest, ) -> tuple[ DatasourceRecord | None, DatasourceGovernance, tuple[DatasourceField, ...] | None, ]: if not request.target_datasource_ref: return ( None, request.governance or _default_governance(mode="static", provider_ref=None), None, ) target = _required_datasource( session, tenant_id=tenant_id, datasource_ref=request.target_datasource_ref, for_update=True, ) if target.mode == "live" or target.shape != "tabular": raise DatasourceValidationError( "Produced rows require a static or cached tabular datasource." ) return ( target, request.governance or _datasource_governance(target), _fields(target.schema_) if target.current_materialization_id else None, ) def _publication_provenance( request: DatasourcePublicationRequest, prepared: _PreparedPublication, ) -> dict[str, object]: return { **dict(request.provenance), "producer_module": prepared.producer_module, "producer_run_ref": prepared.producer_run_ref, "idempotency_key": prepared.idempotency_key, "published_at": utcnow().isoformat(), } def _create_publication_record( session: Session, *, tenant_id: str, actor_id: str | None, datasource: DatasourceRecord, materialization: DatasourceMaterializationRecord, request: DatasourcePublicationRequest, prepared: _PreparedPublication, ) -> DatasourcePublicationRecord: publication = DatasourcePublicationRecord( tenant_id=tenant_id, producer_module=prepared.producer_module, producer_run_ref=prepared.producer_run_ref, idempotency_key=prepared.idempotency_key, request_hash=prepared.request_hash, datasource_id=datasource.id, materialization_id=materialization.id, status="published", details_={ "fingerprint": prepared.fingerprint, "row_count": len(prepared.rows), "set_current": request.set_current, "frozen": request.freeze, "validation": materialization.provenance_.get( "publication_validation", {}, ), }, created_by=actor_id, ) session.add(publication) session.flush() return publication def _publication_request_hash( request: DatasourcePublicationRequest, *, normalized: Sequence[Mapping[str, object]], fingerprint: str, ) -> str: payload = { "producer_module": request.producer_module.strip(), "producer_run_ref": request.producer_run_ref.strip(), "target_datasource_ref": request.target_datasource_ref, "name": request.name, "source_name": request.source_name, "description": request.description, "rows": [dict(row) for row in normalized], "fingerprint": fingerprint, "freeze": request.freeze, "frozen_label": request.frozen_label, "set_current": request.set_current, "source_timestamp": request.source_timestamp, "provenance": dict(request.provenance), "metadata": dict(request.metadata), } encoded = json.dumps( payload, sort_keys=True, separators=(",", ":"), default=str, ) return hashlib.sha256(encoded.encode("utf-8")).hexdigest() __all__ = [ "ADMIN_SCOPE", "CATALOGUE_READ_SCOPE", "SOURCE_WRITE_SCOPE", "STAGE_WRITE_SCOPE", "SqlDatasourceProvider", ]