DataOps-platform/app/runner/api.py

"""Minimal HTTP surface for the independent DataOps Runner."""

import hashlib
import json

from flask import Flask, jsonify, request

from app.runner.auth import TaskTokenInvalid
from app.runner.nodes import NodeExecutionError


def create_runner_app(*, verifier, ledger, registry):
    app = Flask("dataops-runner")
    app.config["MAX_CONTENT_LENGTH"] = 2 * 1024 * 1024

    @app.get("/health")
    def health():
        return jsonify({"status": "ok"})

    def finish_safely(jti, **outcome):
        try:
            ledger.finish(jti, **outcome)
            return True
        except Exception:
            app.logger.error("runner task ledger unavailable")
            return False

    @app.post("/v1/tasks/execute")
    def execute_task():
        payload = request.get_json(silent=True)
        if not isinstance(payload, dict):
            return jsonify({"error": "invalid task request"}), 400
        node = payload.get("node")
        parameters = payload.get("parameters", {})
        if not isinstance(node, dict) or not isinstance(parameters, dict):
            return jsonify({"error": "invalid task request"}), 400
        try:
            claims = verifier.verify(payload.get("task_token"), node=node)
        except TaskTokenInvalid:
            return jsonify({"error": "task token is invalid"}), 401
        binding = {
            "task_uid": claims.task_uid,
            "dataflow_uid": claims.dataflow_uid,
            "deployment_id": claims.deployment_id,
            "environment": claims.environment,
            "workflow_version": claims.workflow_version,
            "correlation_id": claims.correlation_id,
            "node_id": claims.node_id,
            "node_type": claims.node_type,
            "data_source_uid": node.get("data_source_uid"),
            "idempotency_key": (node.get("idempotency") or {}).get("key"),
        }
        try:
            claimed = ledger.claim(
                claims.jti,
                binding,
                expires_at=claims.expires_at,
            )
        except Exception:
            app.logger.error("runner task ledger unavailable")
            return jsonify({"error": "task ledger is unavailable"}), 503
        if not claimed:
            try:
                existing = ledger.get(claims.jti)
            except Exception:
                return jsonify({"error": "task ledger is unavailable"}), 503
            if (
                existing is None
                or any(
                    existing.binding.get(key) != value
                    for key, value in binding.items()
                )
            ):
                return jsonify({"error": "task token already consumed"}), 409
            if existing.status == "running":
                recovered = None
                if node.get("type") in {"rule.apply", "quality.check"}:
                    try:
                        recovered = registry.replay_task(
                            node,
                            correlation_id=claims.correlation_id,
                            deployment_id=claims.deployment_id,
                            task_jti=claims.jti,
                        )
                    except NodeExecutionError:
                        return jsonify(
                            {"error": "task ledger is unavailable"}
                        ), 503
                if isinstance(recovered, dict):
                    replay_body = {
                        "task_uid": claims.task_uid,
                        "correlation_id": claims.correlation_id,
                        "result": recovered,
                    }
                    output_artifact = recovered.get("output_artifact")
                    if isinstance(output_artifact, str):
                        replay_body["output_artifact"] = output_artifact
                    if not finish_safely(
                        claims.jti,
                        status="success",
                        commit_outcome=recovered.get(
                            "commit_outcome",
                            "not_applicable",
                        ),
                        safe_detail="task response recovered",
                        replay_http_status=200,
                        replay_body=replay_body,
                    ):
                        return jsonify(
                            {"error": "task ledger is unavailable"}
                        ), 503
                    response = jsonify(replay_body)
                    response.headers["X-Idempotent-Replay"] = "true"
                    return response
                response = jsonify(
                    {
                        "task_uid": claims.task_uid,
                        "correlation_id": claims.correlation_id,
                        "status": "running",
                    }
                )
                response.headers["Retry-After"] = "2"
                return response, 202
            if (
                existing.replay_body is None
                or existing.replay_http_status is None
                or existing.replay_digest is None
            ):
                return jsonify({"error": "task token already consumed"}), 409
            encoded = json.dumps(
                existing.replay_body,
                sort_keys=True,
                separators=(",", ":"),
                ensure_ascii=False,
            ).encode("utf-8")
            if (
                hashlib.sha256(encoded).hexdigest()
                != existing.replay_digest
            ):
                return jsonify({"error": "task ledger is unavailable"}), 503
            response = jsonify(existing.replay_body)
            response.headers["X-Idempotent-Replay"] = "true"
            return response, existing.replay_http_status
        try:
            result = registry.execute(
                node,
                parameters,
                write_authorized=claims.write_authorized,
                correlation_id=claims.correlation_id,
                dataflow_uid=claims.dataflow_uid,
                deployment_id=claims.deployment_id,
                environment=claims.environment,
                workflow_version=claims.workflow_version,
                node_id=claims.node_id,
                task_jti=claims.jti,
            )
        except NodeExecutionError as exc:
            recorded = finish_safely(
                claims.jti,
                status=(
                    "unknown"
                    if exc.commit_outcome == "unknown"
                    else "failed"
                ),
                commit_outcome=exc.commit_outcome,
                safe_detail=str(exc),
            )
            if not recorded:
                return jsonify({"error": "task ledger is unavailable"}), 503
            return jsonify({"error": str(exc)}), 400
        except Exception:
            finish_safely(
                claims.jti,
                status="failed",
                safe_detail="task execution failed",
            )
            app.logger.exception("runner task failed")
            return jsonify({"error": "task execution failed"}), 500
        commit_outcome = (
            result.get("commit_outcome", "not_applicable")
            if isinstance(result, dict)
            else "not_applicable"
        )
        response = {
            "task_uid": claims.task_uid,
            "correlation_id": claims.correlation_id,
            "result": result,
        }
        if (
            isinstance(result, dict)
            and isinstance(result.get("output_artifact"), str)
        ):
            response["output_artifact"] = result["output_artifact"]
        replay_body = (
            response
            if node.get("type") in {"rule.apply", "quality.check"}
            else None
        )
        if not finish_safely(
            claims.jti,
            status="success",
            commit_outcome=commit_outcome,
            safe_detail="task completed",
            replay_http_status=(200 if replay_body is not None else None),
            replay_body=replay_body,
        ):
            return jsonify({"error": "task ledger is unavailable"}), 503
        return jsonify(response)

    return app