DataOps-platform/app/runner/rule_polars.py

"""Runner adapter for OS-isolated, allowlisted Polars execution."""

from __future__ import annotations

import os
import tempfile
from contextlib import ExitStack, suppress
from typing import Any

from app.core.common.identifiers import ensure_governance_uid
from app.core.data_rules.compilers.polars import (
    bound_polars_plan_hash,
    validate_bound_polars_plan,
)
from app.runner.artifacts import ArtifactCommitUnknown
from app.runner.nodes import NodeExecutionError
from app.runner.polars_worker import (
    PolarsWorkerError,
    PolarsWorkerResourceError,
    execute_isolated_polars_plan,
)

_IDEMPOTENCY = {
    "deduplication_key",
    "partition_replace",
    "upsert",
}
_PUBLIC_ARTIFACT_KEYS = (
    "artifact_ref",
    "digest",
    "row_count",
    "schema_hash",
    "expires_at",
)
_RESULT_KEYS = (
    "rows_in",
    "rows_out",
    "rows_rejected",
    "rows_filtered",
    "rows_deduplicated",
    "rows_join_dropped",
    "rows_aggregated",
    "violation_count",
    "violations",
    "_violation_sample",
)


def _uid(value: Any, label: str) -> str:
    try:
        return ensure_governance_uid({"uid": str(value)})
    except ValueError as exc:
        raise NodeExecutionError(f"{label} is invalid") from exc


def _resolve_artifact(
    resolver,
    *,
    binding_id: str,
    binding_hash: str,
    correlation_id: str,
    kind: str,
) -> dict[str, Any]:
    try:
        artifact = resolver.resolve(
            binding_id=binding_id,
            correlation_id=correlation_id,
            kind=kind,
        )
    except Exception as exc:
        raise NodeExecutionError(
            "published Polars artifact binding was not resolved"
        ) from exc
    if (
        not isinstance(artifact, dict)
        or artifact.get("binding_hash") != binding_hash
        or not isinstance(artifact.get("artifact_ref"), str)
        or not isinstance(artifact.get("digest"), str)
        or not isinstance(artifact.get("schema_fields"), list)
    ):
        raise NodeExecutionError(
            "published Polars artifact binding does not match"
        )
    return artifact


class PolarsRulePlanAdapter:
    """Stage bounded inputs and execute the plan in a hard-limited process."""

    def __init__(
        self,
        *,
        artifact_store,
        artifact_resolver,
        masking_policies=None,
        artifact_ttl_seconds=3600,
    ):
        self.artifact_store = artifact_store
        self.artifact_resolver = artifact_resolver
        self.masking_policies = dict(masking_policies or {})
        self.artifact_ttl_seconds = int(artifact_ttl_seconds)

    def execute(
        self,
        *,
        plan,
        node,
        parameters,
        write_authorized,
        correlation_id=None,
    ):
        try:
            normalized = validate_bound_polars_plan(plan)
        except ValueError as exc:
            raise NodeExecutionError(
                "published Polars rule plan is invalid"
            ) from exc
        config = node.get("config") or {}
        if config.get("execution_plan_hash") != bound_polars_plan_hash(
            normalized
        ):
            raise NodeExecutionError(
                "published Polars rule plan hash does not match"
            )
        if config.get("rule_version_id") != normalized["rule_version_id"]:
            raise NodeExecutionError(
                "published Polars rule id does not match"
            )
        idempotency = node.get("idempotency")
        if (
            node.get("type") != "rule.apply"
            or node.get("purpose") != "write"
            or not write_authorized
            or not isinstance(idempotency, dict)
            or idempotency.get("strategy") not in _IDEMPOTENCY
            or not str(idempotency.get("key") or "").strip()
        ):
            raise NodeExecutionError(
                "governed write authorization and idempotency are required"
            )
        if parameters in ({}, None):
            input_handoff = None
        elif (
            isinstance(parameters, dict)
            and set(parameters) == {"input_artifact"}
            and isinstance(parameters["input_artifact"], str)
        ):
            input_handoff = parameters["input_artifact"]
        else:
            raise NodeExecutionError(
                "bound Polars rule plans accept only one artifact handoff"
            )
        correlation = _uid(correlation_id, "correlation_id")
        try:
            self.artifact_resolver.attest_binding(
                binding_id=normalized["output_binding_id"],
                binding_hash=normalized["output_binding_hash"],
                access_mode="write",
            )
        except Exception as exc:
            raise NodeExecutionError(
                "published Polars output binding no longer matches"
            ) from exc
        if input_handoff is None:
            source = _resolve_artifact(
                self.artifact_resolver,
                binding_id=normalized["input_binding_id"],
                binding_hash=normalized["input_binding_hash"],
                correlation_id=correlation,
                kind="input",
            )
        else:
            try:
                self.artifact_resolver.attest_binding(
                    binding_id=normalized["input_binding_id"],
                    binding_hash=normalized["input_binding_hash"],
                    access_mode="read",
                )
                source = self.artifact_resolver.resolve_handoff(
                    artifact_ref=input_handoff,
                    correlation_id=correlation,
                )
            except Exception as exc:
                raise NodeExecutionError(
                    "upstream Polars artifact handoff was not resolved"
                ) from exc
        limits = normalized["resource_limits"]
        output_path = None
        try:
            with ExitStack() as stack:
                try:
                    input_path = stack.enter_context(
                        self.artifact_store.stage(
                            source["artifact_ref"],
                            source["digest"],
                            expected_schema_fields=normalized[
                                "input_fields"
                            ],
                            limits=limits,
                        )
                    )
                    lookup_paths = {}
                    for operation in normalized["operations"]:
                        if operation["op"] != "lookup_join":
                            continue
                        lookup_artifact = _resolve_artifact(
                            self.artifact_resolver,
                            binding_id=operation[
                                "lookup_binding_id"
                            ],
                            binding_hash=operation[
                                "lookup_binding_hash"
                            ],
                            correlation_id=correlation,
                            kind="lookup",
                        )
                        lookup_paths[
                            operation["lookup_binding_id"]
                        ] = stack.enter_context(
                            self.artifact_store.stage(
                                lookup_artifact["artifact_ref"],
                                lookup_artifact["digest"],
                                expected_schema_fields=operation[
                                    "lookup_fields"
                                ],
                                limits=limits,
                            )
                        )
                except ValueError as exc:
                    raise NodeExecutionError(
                        "published Polars input artifact is invalid"
                    ) from exc
                with tempfile.NamedTemporaryFile(
                    prefix="dataops-polars-worker-output-",
                    suffix=".parquet",
                    delete=False,
                ) as handle:
                    output_path = handle.name
                try:
                    worker_result = execute_isolated_polars_plan(
                        {
                            "plan": normalized,
                            "input_path": input_path,
                            "lookup_paths": lookup_paths,
                            "output_path": output_path,
                            "masking_policies": self.masking_policies,
                        },
                        memory_limit_bytes=limits[
                            "memory_limit_bytes"
                        ],
                    )
                except PolarsWorkerResourceError as exc:
                    raise NodeExecutionError(
                        "published Polars worker exceeded its hard memory limit"
                    ) from exc
                except PolarsWorkerError as exc:
                    raise NodeExecutionError(
                        "published Polars worker execution failed"
                    ) from exc
                try:
                    registered = self.artifact_resolver.publish_path(
                        output_path,
                        binding_id=normalized["output_binding_id"],
                        binding_hash=normalized["output_binding_hash"],
                        correlation_id=correlation,
                        kind="output",
                        ttl_seconds=self.artifact_ttl_seconds,
                        schema_fields=normalized["output_fields"],
                        limits=limits,
                    )
                except ArtifactCommitUnknown as exc:
                    raise NodeExecutionError(
                        "published Polars artifact commit outcome is unknown",
                        commit_outcome="unknown",
                    ) from exc
                except Exception as exc:
                    raise NodeExecutionError(
                        "published Polars output artifact publication failed"
                    ) from exc
        finally:
            if output_path is not None:
                with suppress(FileNotFoundError):
                    os.unlink(output_path)
        return {
            **{key: registered[key] for key in _PUBLIC_ARTIFACT_KEYS},
            **{key: worker_result[key] for key in _RESULT_KEYS},
            "commit_outcome": "committed",
        }