from __future__ import annotations import hashlib import json from govoplan_core.auth import ApiPrincipal from govoplan_core.core.dataflows import ( DataflowDatasetDescriptor, DataflowDatasetRequest, DataflowDatasetResult, DataflowRunConflictError, DataflowRunUnavailableError, ) from govoplan_core.core.datasources import ( DatasourceError, DatasourceReadRequest, datasource_catalogue, ) from govoplan_core.security.time import utc_now from govoplan_dataflow.backend.backends.base import ExecutionBudget from govoplan_dataflow.backend.executor import PipelineExecutionError from govoplan_dataflow.backend.governance import require_definition_action from govoplan_dataflow.backend.schemas import PipelineGraph from govoplan_dataflow.backend.service import ( _execute_pipeline_preview, get_pipeline, get_pipeline_revision, get_pipeline_run, list_pipelines, ) from govoplan_dataflow.backend.subflows import substitute_parameters class SqlDataflowDatasetOutputProvider: def __init__(self, registry: object | None = None) -> None: self.registry = registry def list_outputs( self, session: object, principal: object, *, query: str = "", limit: int = 100, ) -> tuple[DataflowDatasetDescriptor, ...]: typed_session, typed_principal = _contracts(session, principal) normalized_query = query.strip().casefold() result: list[DataflowDatasetDescriptor] = [] for pipeline in list_pipelines( typed_session, tenant_id=typed_principal.tenant_id, ): if normalized_query and normalized_query not in pipeline.name.casefold(): continue try: decision = require_definition_action( pipeline, principal=typed_principal, registry=self.registry, action="run", ) except PermissionError: continue revision = get_pipeline_revision(typed_session, pipeline=pipeline) result.append( DataflowDatasetDescriptor( pipeline_ref=pipeline.id, name=pipeline.name, description=pipeline.description, revision=revision.revision, definition_hash=revision.content_hash, status=pipeline.status, updated_at=pipeline.updated_at, parameters=dict(pipeline.metadata_.get("parameters") or {}), provenance={ "module": "dataflow", "scope_type": pipeline.scope_type, "scope_id": pipeline.scope_id, "policy_decision": decision.to_dict(), }, ) ) if len(result) >= max(1, min(limit, 200)): break return tuple(result) def read_output( self, session: object, principal: object, *, request: DataflowDatasetRequest, ) -> DataflowDatasetResult: typed_session, typed_principal = _contracts(session, principal) row_limit = max(1, min(request.row_limit, 2_000)) pipeline = get_pipeline( typed_session, tenant_id=typed_principal.tenant_id, pipeline_id=request.pipeline_ref, ) decision = require_definition_action( pipeline, principal=typed_principal, registry=self.registry, action="run", ) revision = get_pipeline_revision( typed_session, pipeline=pipeline, revision=request.revision, ) if ( request.expected_definition_hash and request.expected_definition_hash != revision.content_hash ): raise DataflowRunConflictError( "The pinned Dataflow definition hash no longer matches the requested revision." ) if request.run_ref: return _read_published_run_output( typed_session, typed_principal, registry=self.registry, request=request, pipeline=pipeline, revision=revision, policy_decision=decision.to_dict(), row_limit=row_limit, ) graph = PipelineGraph.model_validate( substitute_parameters(revision.graph, dict(request.parameters)) ) try: result, executor_version = _execute_pipeline_preview( graph, session=typed_session, principal=typed_principal, registry=self.registry, backend="auto", row_limit=row_limit, preview_node_id=None, budget=ExecutionBudget( max_output_rows=row_limit, max_batch_bytes=4_000_000, max_wall_seconds=5.0, ), ) except PipelineExecutionError as exc: raise DataflowRunUnavailableError( f"The pinned Dataflow output could not be evaluated: {exc}" ) from exc source_fingerprints = tuple( dict(item) for item in result.source_fingerprints ) if request.expected_source_fingerprints and not _fingerprints_match( request.expected_source_fingerprints, source_fingerprints, ): raise DataflowRunConflictError( "Dataflow source fingerprints differ from the pinned source state." ) output_hash = hashlib.sha256( json.dumps( result.rows, sort_keys=True, separators=(",", ":"), ensure_ascii=True, default=str, ).encode("utf-8") ).hexdigest() return DataflowDatasetResult( pipeline_ref=pipeline.id, revision=revision.revision, definition_hash=revision.content_hash, rows=tuple(dict(item) for item in result.rows), total_rows=result.total_rows, truncated=result.truncated, output_hash=output_hash, executor_version=executor_version, source_fingerprints=source_fingerprints, diagnostics=tuple( item.model_dump(mode="json") for item in result.diagnostics ), generated_at=utc_now(), provenance={ "module": "dataflow", "scope_type": pipeline.scope_type, "scope_id": pipeline.scope_id, "policy_decision": decision.to_dict(), "parameters": dict(request.parameters), }, ) def dataset_output_provider(context: object | None = None): return SqlDataflowDatasetOutputProvider(getattr(context, "registry", None)) def _read_published_run_output( session, principal: ApiPrincipal, *, registry: object | None, request: DataflowDatasetRequest, pipeline, revision, policy_decision: dict[str, object], row_limit: int, ) -> DataflowDatasetResult: if request.parameters: raise DataflowRunConflictError( "An immutable published run cannot be evaluated with new parameters." ) run = get_pipeline_run( session, tenant_id=principal.tenant_id, run_ref=request.run_ref or "", ) if run.pipeline_id != pipeline.id or run.pipeline_revision_id != revision.id: raise DataflowRunConflictError( "The published run does not belong to the pinned Dataflow revision." ) if run.definition_hash != revision.content_hash: raise DataflowRunConflictError( "The published run definition evidence does not match the pinned revision." ) if run.status != "succeeded": raise DataflowRunUnavailableError( "Only a successful Dataflow run can be used as an immutable dataset." ) if not run.output_datasource_ref or not run.output_materialization_ref: raise DataflowRunUnavailableError( "The successful Dataflow run has no immutable Datasource publication." ) provider = datasource_catalogue(registry) if provider is None: raise DataflowRunUnavailableError( "The Datasource catalogue required by this published run is not enabled." ) rows: list[dict[str, object]] = [] materialization = None total_rows = 0 try: while len(rows) < row_limit: remaining = row_limit - len(rows) page = provider.read_datasource( session, principal, request=DatasourceReadRequest( datasource_ref=run.output_datasource_ref, materialization_ref=run.output_materialization_ref, limit=min(500, remaining), offset=len(rows), expected_fingerprint=( materialization.fingerprint if materialization is not None else None ), ), ) if ( page.materialization is None or page.materialization.ref != run.output_materialization_ref ): raise DataflowRunConflictError( "The Datasource provider returned a different output materialization." ) if len(page.rows) > remaining: raise DataflowRunConflictError( "The Datasource provider exceeded the requested output window." ) materialization = page.materialization total_rows = page.total_rows rows.extend(dict(item) for item in page.rows) if not page.truncated: break if not page.rows: raise DataflowRunUnavailableError( "The Datasource provider made no progress while reading the published output." ) except DatasourceError as exc: raise DataflowRunUnavailableError( "The immutable Datasource output is unavailable to the current principal." ) from exc source_fingerprints = tuple(dict(item) for item in run.source_fingerprints) if request.expected_source_fingerprints and not _fingerprints_match( request.expected_source_fingerprints, source_fingerprints, ): raise DataflowRunConflictError( "Dataflow source fingerprints differ from the pinned run evidence." ) output_hash = hashlib.sha256( json.dumps( rows, sort_keys=True, separators=(",", ":"), ensure_ascii=True, default=str, ).encode("utf-8") ).hexdigest() if materialization is None: raise DataflowRunUnavailableError( "The Datasource provider returned no materialization evidence." ) return DataflowDatasetResult( pipeline_ref=pipeline.id, revision=revision.revision, definition_hash=revision.content_hash, rows=tuple(rows), total_rows=total_rows, truncated=len(rows) < total_rows, output_hash=output_hash, executor_version=run.executor_version, run_ref=f"dataflow-run:{run.id}", source_fingerprints=source_fingerprints, diagnostics=tuple(dict(item) for item in run.diagnostics), generated_at=run.finished_at or materialization.created_at, provenance={ "module": "dataflow", "scope_type": pipeline.scope_type, "scope_id": pipeline.scope_id, "policy_decision": policy_decision, "immutable_run": True, "publication_ref": run.output_publication_ref, "datasource_ref": run.output_datasource_ref, "materialization_ref": run.output_materialization_ref, "materialization_fingerprint": materialization.fingerprint, "governance": materialization.governance.to_dict(), }, ) def _fingerprints_match(expected, actual) -> bool: def normalized(values): return sorted( json.dumps(dict(item), sort_keys=True, separators=(",", ":"), default=str) for item in values ) return normalized(expected) == normalized(actual) def _contracts(session: object, principal: object): if not hasattr(session, "scalar") or not hasattr(session, "scalars"): raise TypeError("Dataflow dataset output requires a SQLAlchemy session.") if not isinstance(principal, ApiPrincipal): raise TypeError("Dataflow dataset output requires an API principal.") return session, principal __all__ = ["SqlDataflowDatasetOutputProvider", "dataset_output_provider"]