"""Single-use task-token ledger interfaces and an in-memory test implementation.""" from __future__ import annotations import hashlib import json import threading import time 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 lease_expires_at: float = 0 class InMemoryTaskLedger: def __init__(self, *, clock=None, lease_seconds=300): self._records = {} self._lock = threading.Lock() self._clock = clock or time.time self._lease_seconds = int(lease_seconds) 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), lease_expires_at=( float(self._clock()) + self._lease_seconds ), ) 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)) def reconcile_running(self, jti, binding): with self._lock: record = self._records.get(str(jti)) if ( record is not None and record.status == "running" and record.binding == dict(binding) and record.lease_expires_at <= float(self._clock()) ): record.status = "unknown" record.commit_outcome = "unknown" record.safe_detail = "task execution lease expired" return record class PostgresTaskLedger: """Durable, cross-worker single-use ledger backed by the platform DB.""" def __init__(self, engine, *, lease_seconds=300): self.engine = engine self.lease_seconds = int(lease_seconds) if self.lease_seconds < 30 or self.lease_seconds > 900: raise ValueError("runner task lease is invalid") 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), "lease_seconds": self.lease_seconds, } 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, lease_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), CAST(:deployment_id AS uuid), :environment, CURRENT_TIMESTAMP + make_interval(secs => :lease_seconds) ) ON CONFLICT DO NOTHING """ ), parameters, ) return int(result.rowcount or 0) == 1 def reconcile_running(self, jti, binding): 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"], } with self.engine.begin() as connection: connection.execute( text( """ UPDATE public.runner_task_executions SET status = 'unknown', commit_outcome = 'unknown', safe_detail = 'task execution lease expired', finished_at = CURRENT_TIMESTAMP WHERE token_jti = CAST(:jti AS uuid) AND task_uid = CAST(:task_uid AS uuid) AND dataflow_uid = CAST(:dataflow_uid AS uuid) AND deployment_id = CAST(:deployment_id AS uuid) AND environment = :environment AND workflow_version = :workflow_version AND correlation_id = CAST(:correlation_id AS uuid) AND node_id = :node_id AND node_type = :node_type AND status = 'running' AND lease_expires_at <= CURRENT_TIMESTAMP """ ), parameters, ) return self.get(jti) 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, EXTRACT( EPOCH FROM lease_expires_at )::double precision AS lease_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"], lease_expires_at=float(row["lease_expires_at"]), )