from __future__ import annotations from collections.abc import Mapping, Sequence from urllib.parse import quote from sqlalchemy import func, or_, select from sqlalchemy.orm import Session from govoplan_core.auth import ApiPrincipal from govoplan_core.core.modules import ModuleContext from govoplan_core.core.search import ( SearchAuthorizationRequest, SearchBackfillPage, SearchBackfillRequest, SearchDocument, SearchResourceType, ) from govoplan_dataflow.backend.db.models import DataflowPipeline from govoplan_dataflow.backend.governance import definition_decision PROVIDER_ID = "dataflow.pipelines" RESOURCE_TYPE = "dataflow_pipeline" READ_SCOPE = "dataflow:pipeline:read" ADMIN_SCOPE = "dataflow:pipeline:admin" class DataflowSearchSource: def __init__(self, registry: object | None) -> None: self.registry = registry def resource_types(self) -> Sequence[SearchResourceType]: return ( SearchResourceType( provider_id=PROVIDER_ID, module_id="dataflow", resource_type=RESOURCE_TYPE, label="Data pipelines", requires_authorization_recheck=True, ), ) def backfill( self, session: object, *, request: SearchBackfillRequest, ) -> SearchBackfillPage: db = _session(session) if ( request.provider_id != PROVIDER_ID or request.resource_type != RESOURCE_TYPE ): raise ValueError("Unsupported Dataflow search source.") source_filter = or_( DataflowPipeline.tenant_id == request.tenant_id, DataflowPipeline.tenant_id.is_(None), ) statement = select(DataflowPipeline).where( source_filter, DataflowPipeline.deleted_at.is_(None), ) if request.cursor: statement = statement.where( DataflowPipeline.id > request.cursor ) rows = list( db.scalars( statement.order_by(DataflowPipeline.id).limit( request.limit + 1 ) ) ) has_more = len(rows) > request.limit selected = rows[: request.limit] high_watermark = db.scalar( select(func.max(DataflowPipeline.updated_at)).where( source_filter, DataflowPipeline.deleted_at.is_(None), ) ) return SearchBackfillPage( documents=tuple( _pipeline_document( pipeline, tenant_id=request.tenant_id, ) for pipeline in selected ), next_cursor=selected[-1].id if has_more and selected else None, complete=not has_more, high_watermark=( high_watermark.isoformat() if high_watermark is not None else None ), ) def authorize( self, session: object, principal: object, *, requests: Sequence[SearchAuthorizationRequest], ) -> Mapping[str, bool]: decisions = { request.reference.key: False for request in requests } if not isinstance(principal, ApiPrincipal): return decisions if not ( principal.has(READ_SCOPE) or principal.has(ADMIN_SCOPE) ): return decisions valid_requests = tuple( request for request in requests if ( request.reference.tenant_id == principal.tenant_id and request.reference.module_id == "dataflow" and request.reference.resource_type == RESOURCE_TYPE ) ) if not valid_requests: return decisions db = _session(session) ids = { request.reference.resource_id for request in valid_requests } pipelines = { pipeline.id: pipeline for pipeline in db.scalars( select(DataflowPipeline).where( DataflowPipeline.id.in_(ids), or_( DataflowPipeline.tenant_id == principal.tenant_id, DataflowPipeline.tenant_id.is_(None), ), DataflowPipeline.deleted_at.is_(None), ) ) } for request in valid_requests: pipeline = pipelines.get( request.reference.resource_id ) if pipeline is None: continue try: allowed = definition_decision( pipeline, principal=principal, registry=self.registry, action="view", ).allowed except (PermissionError, RuntimeError, ValueError): allowed = False decisions[request.reference.key] = allowed return decisions def create_dataflow_search_source( context: ModuleContext, ) -> DataflowSearchSource: return DataflowSearchSource(context.registry) def _pipeline_document( pipeline: DataflowPipeline, *, tenant_id: str, ) -> SearchDocument: updated_at = pipeline.updated_at source_revision = ( f"{pipeline.current_revision}:" f"{updated_at.isoformat() if updated_at is not None else 'unknown'}" ) return SearchDocument( tenant_id=tenant_id, module_id="dataflow", provider_id=PROVIDER_ID, resource_type=RESOURCE_TYPE, resource_id=pipeline.id, title=pipeline.name, url=f"/dataflow?pipelineId={quote(pipeline.id, safe='')}", summary=pipeline.description, body=pipeline.description, keywords=( pipeline.status, pipeline.scope_type, pipeline.definition_kind, ), visibility="tenant", source_revision=source_revision, source_updated_at=updated_at, metadata={ "status": pipeline.status, "scope_type": pipeline.scope_type, "definition_kind": pipeline.definition_kind, "current_revision": pipeline.current_revision, }, requires_authorization_recheck=True, ) def _session(value: object) -> Session: if not isinstance(value, Session): raise TypeError( "Dataflow search requires a SQLAlchemy session." ) return value __all__ = [ "DataflowSearchSource", "PROVIDER_ID", "RESOURCE_TYPE", "create_dataflow_search_source", ]