DataOps-platform/app/runner/ledger.py

"""Single-use task-token ledger interfaces and an in-memory test implementation."""

from __future__ import annotations

import hashlib
import json
import threading
from collections.abc import Mapping
from dataclasses import dataclass

from sqlalchemy import text


@dataclass
class TaskLedgerRecord:
    jti: str
    binding: Mapping[str, object]
    expires_at: int
    status: str = "running"
    commit_outcome: str = "not_applicable"
    safe_detail: str = ""
    replay_http_status: int | None = None
    replay_body: Mapping[str, object] | None = None
    replay_digest: str | None = None


class InMemoryTaskLedger:
    def __init__(self):
        self._records = {}
        self._lock = threading.Lock()

    def claim(self, jti, binding, *, expires_at):
        with self._lock:
            if jti in self._records:
                return False
            self._records[jti] = TaskLedgerRecord(
                jti=str(jti),
                binding=dict(binding),
                expires_at=int(expires_at),
            )
            return True

    def finish(
        self,
        jti,
        *,
        status,
        commit_outcome="not_applicable",
        safe_detail="",
        replay_http_status=None,
        replay_body=None,
    ):
        with self._lock:
            record = self._records[str(jti)]
            record.status = str(status)
            record.commit_outcome = str(commit_outcome)
            record.safe_detail = str(safe_detail)[:500]
            if replay_body is not None:
                encoded = json.dumps(
                    replay_body,
                    sort_keys=True,
                    separators=(",", ":"),
                    ensure_ascii=False,
                ).encode("utf-8")
                record.replay_http_status = int(replay_http_status)
                record.replay_body = dict(replay_body)
                record.replay_digest = hashlib.sha256(encoded).hexdigest()

    def get(self, jti) -> TaskLedgerRecord | None:
        with self._lock:
            return self._records.get(str(jti))


class PostgresTaskLedger:
    """Durable, cross-worker single-use ledger backed by the platform DB."""

    def __init__(self, engine):
        self.engine = engine

    def claim(self, jti, binding, *, expires_at):
        parameters = {
            "jti": str(jti),
            "task_uid": binding["task_uid"],
            "dataflow_uid": binding["dataflow_uid"],
            "deployment_id": binding["deployment_id"],
            "environment": binding["environment"],
            "workflow_version": int(binding["workflow_version"]),
            "correlation_id": binding["correlation_id"],
            "node_id": binding["node_id"],
            "node_type": binding["node_type"],
            "data_source_uid": binding.get("data_source_uid"),
            "idempotency_key": binding.get("idempotency_key"),
            "expires_at": int(expires_at),
        }
        with self.engine.begin() as connection:
            result = connection.execute(
                text(
                    """
                    INSERT INTO public.runner_task_executions (
                        token_jti, task_uid, dataflow_uid, workflow_version,
                        correlation_id, node_id, node_type, data_source_uid,
                        idempotency_key, status, commit_outcome, expires_at,
                        deployment_id, environment
                    ) VALUES (
                        CAST(:jti AS uuid), CAST(:task_uid AS uuid),
                        CAST(:dataflow_uid AS uuid), :workflow_version,
                        CAST(:correlation_id AS uuid), :node_id, :node_type,
                        CAST(:data_source_uid AS uuid), :idempotency_key,
                        'running', 'not_applicable',
                        to_timestamp(:expires_at),
                        CAST(:deployment_id AS uuid), :environment
                    )
                    ON CONFLICT DO NOTHING
                    """
                ),
                parameters,
            )
        return int(result.rowcount or 0) == 1

    def finish(
        self,
        jti,
        *,
        status,
        commit_outcome="not_applicable",
        safe_detail="",
        replay_http_status=None,
        replay_body=None,
    ):
        allowed_statuses = {"success", "failed", "unknown"}
        allowed_outcomes = {
            "not_applicable",
            "not_committed",
            "committed",
            "unknown",
        }
        if status not in allowed_statuses or commit_outcome not in allowed_outcomes:
            raise ValueError("runner task outcome is invalid")
        replay_digest = None
        if replay_body is not None:
            encoded = json.dumps(
                replay_body,
                sort_keys=True,
                separators=(",", ":"),
                ensure_ascii=False,
            ).encode("utf-8")
            if len(encoded) > 32_768:
                raise ValueError("runner replay body exceeds the safe limit")
            if replay_http_status != 200:
                raise ValueError("runner replay status is invalid")
            replay_digest = hashlib.sha256(encoded).hexdigest()
        with self.engine.begin() as connection:
            connection.execute(
                text(
                    """
                    UPDATE public.runner_task_executions
                    SET status = :status,
                        commit_outcome = :commit_outcome,
                        safe_detail = :safe_detail,
                        replay_http_status = :replay_http_status,
                        replay_body = CAST(:replay_body AS jsonb),
                        replay_digest = :replay_digest,
                        finished_at = CURRENT_TIMESTAMP
                    WHERE token_jti = CAST(:jti AS uuid)
                      AND status = 'running'
                    """
                ),
                {
                    "jti": str(jti),
                    "status": status,
                    "commit_outcome": commit_outcome,
                    "safe_detail": str(safe_detail)[:500],
                    "replay_http_status": replay_http_status,
                    "replay_body": (
                        json.dumps(replay_body)
                        if replay_body is not None
                        else None
                    ),
                    "replay_digest": replay_digest,
                },
            )

    def get(self, jti):
        with self.engine.begin() as connection:
            row = (
                connection.execute(
                    text(
                        """
                        SELECT token_jti::text AS jti, task_uid::text,
                               dataflow_uid::text, deployment_id::text,
                               environment, workflow_version,
                               correlation_id::text,
                               node_id, data_source_uid::text,
                               idempotency_key, status, commit_outcome,
                               node_type,
                               safe_detail, EXTRACT(EPOCH FROM expires_at)::bigint
                                   AS expires_at,
                               replay_http_status, replay_body,
                               replay_digest
                        FROM public.runner_task_executions
                        WHERE token_jti = CAST(:jti AS uuid)
                        """
                    ),
                    {"jti": str(jti)},
                )
                .mappings()
                .one_or_none()
            )
        if row is None:
            return None
        return TaskLedgerRecord(
            jti=row["jti"],
            binding={
                "task_uid": row["task_uid"],
                "dataflow_uid": row["dataflow_uid"],
                "deployment_id": row["deployment_id"],
                "environment": row["environment"],
                "workflow_version": row["workflow_version"],
                "correlation_id": row["correlation_id"],
                "node_id": row["node_id"],
                "node_type": row["node_type"],
                "data_source_uid": row["data_source_uid"],
                "idempotency_key": row["idempotency_key"],
            },
            expires_at=int(row["expires_at"]),
            status=row["status"],
            commit_outcome=row["commit_outcome"],
            safe_detail=row["safe_detail"] or "",
            replay_http_status=row["replay_http_status"],
            replay_body=(
                dict(row["replay_body"])
                if isinstance(row["replay_body"], dict)
                else row["replay_body"]
            ),
            replay_digest=row["replay_digest"],
        )