DataOps-platform/app/runner/ledger.py

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

from __future__ import annotations

import threading
from dataclasses import dataclass
from typing import Mapping, Optional

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 = ""


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="",
    ):
        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]

    def get(self, jti) -> Optional[TaskLedgerRecord]:
        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"],
            "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
                    ) 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)
                    )
                    ON CONFLICT DO NOTHING
                    """
                ),
                parameters,
            )
        return int(result.rowcount or 0) == 1

    def finish(
        self,
        jti,
        *,
        status,
        commit_outcome="not_applicable",
        safe_detail="",
    ):
        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")
        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,
                        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],
                },
            )

    def get(self, jti):
        with self.engine.begin() as connection:
            row = (
                connection.execute(
                    text(
                        """
                        SELECT token_jti::text AS jti, task_uid::text,
                               node_id, data_source_uid::text,
                               idempotency_key, status, commit_outcome,
                               safe_detail, EXTRACT(EPOCH FROM expires_at)::bigint
                                   AS expires_at
                        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"],
                "node_id": row["node_id"],
                "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 "",
        )