govoplan-dataflow/src/govoplan_dataflow/backend/search_source.py

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",
]