|
| 1 | +"""Immutable heartbeat-resume metadata captured from one negotiated poll.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import json |
| 6 | +import re |
| 7 | +from collections.abc import Mapping |
| 8 | +from dataclasses import dataclass, field |
| 9 | +from datetime import datetime |
| 10 | +from typing import Any |
| 11 | + |
| 12 | +from ._typed_values import decode_payload_value, validate_annotation |
| 13 | +from .errors import ActivityHeartbeatResumeUnavailable |
| 14 | + |
| 15 | +_IDENTITY = ("task_id", "activity_execution_id", "activity_attempt_id", "workflow_id", "run_id", "namespace") |
| 16 | +_TIMESTAMP = re.compile( |
| 17 | + r"[0-9]{4}-[0-9]{2}-[0-9]{2}T(?:[01][0-9]|2[0-3]):[0-5][0-9]:[0-5][0-9]" |
| 18 | + r"(?:\.[0-9]{1,9})?(?:Z|[+-](?:[01][0-9]|2[0-3]):[0-5][0-9])" |
| 19 | +) |
| 20 | + |
| 21 | + |
| 22 | +def _invalid() -> ValueError: |
| 23 | + return ValueError("activity_heartbeat_resume v1 claim has an invalid checkpoint") |
| 24 | + |
| 25 | + |
| 26 | +@dataclass(frozen=True, init=False) |
| 27 | +class ActivityHeartbeatCheckpoint: |
| 28 | + """The acknowledged snapshot received with this claim, never a pending write. |
| 29 | +
|
| 30 | + ``contract_version`` describes the negotiated wire contract. Server v1 |
| 31 | + supplies no checkpoint write sequence. Acknowledgement does not commit an |
| 32 | + external business effect atomically. |
| 33 | + """ |
| 34 | + |
| 35 | + acknowledged_at: str |
| 36 | + contract_version: int = field(default=1, init=False) |
| 37 | + _details_json: str = field(repr=False) |
| 38 | + |
| 39 | + def __init__(self, details: Any, acknowledged_at: str) -> None: |
| 40 | + if details is None or not isinstance(acknowledged_at, str) or not _TIMESTAMP.fullmatch(acknowledged_at): |
| 41 | + raise _invalid() |
| 42 | + try: |
| 43 | + # Python 3.10's parser accepts fewer fractional forms. Validate the |
| 44 | + # calendar/offset separately and retain the original precision. |
| 45 | + calendar = re.sub(r"\.\d+(?=Z|[+-])", "", acknowledged_at) |
| 46 | + datetime.fromisoformat(calendar.replace("Z", "+00:00")) |
| 47 | + encoded = json.dumps(details, allow_nan=False) |
| 48 | + except (ValueError, TypeError, OverflowError, RecursionError): |
| 49 | + raise _invalid() from None |
| 50 | + object.__setattr__(self, "acknowledged_at", acknowledged_at) |
| 51 | + object.__setattr__(self, "_details_json", encoded) |
| 52 | + |
| 53 | + def details(self, *, result_type: Any = None) -> Any: |
| 54 | + """Return an independent copy, optionally decoded using SDK typed values.""" |
| 55 | + value = json.loads(self._details_json) |
| 56 | + if result_type is None: |
| 57 | + return value |
| 58 | + validate_annotation(result_type, "checkpoint") |
| 59 | + return decode_payload_value(value, result_type, "checkpoint") |
| 60 | + |
| 61 | + |
| 62 | +@dataclass(frozen=True) |
| 63 | +class _HeartbeatResume: |
| 64 | + available: bool = False |
| 65 | + checkpoint: ActivityHeartbeatCheckpoint | None = None |
| 66 | + identity: tuple[Any, ...] = () |
| 67 | + |
| 68 | + def get(self) -> ActivityHeartbeatCheckpoint | None: |
| 69 | + if not self.available: |
| 70 | + raise ActivityHeartbeatResumeUnavailable() |
| 71 | + return self.checkpoint |
| 72 | + |
| 73 | + |
| 74 | +class _ActivityTask(dict[str, Any]): |
| 75 | + """Keep poll metadata outside the task's unchanged public wire dictionary.""" |
| 76 | + |
| 77 | + def __init__(self, task: Mapping[str, Any], resume: _HeartbeatResume) -> None: |
| 78 | + super().__init__(task) |
| 79 | + self.heartbeat_resume = resume |
| 80 | + |
| 81 | + |
| 82 | +def _identity(task: Mapping[str, Any]) -> tuple[Any, ...]: |
| 83 | + values = tuple(task.get(name) for name in _IDENTITY) |
| 84 | + attempt = task.get("attempt_number") |
| 85 | + if (any(not isinstance(value, str) or not value.strip() for value in values) |
| 86 | + or type(attempt) is not int or attempt < 1): |
| 87 | + raise _invalid() |
| 88 | + return (*values, attempt, task.get("lease_owner"), task.get("activity_type"), task.get("task_queue")) |
| 89 | + |
| 90 | + |
| 91 | +def capture_poll(data: Mapping[str, Any], namespace: str, worker_id: str, task_queue: str | None = None) -> Any: |
| 92 | + task = data.get("task") |
| 93 | + capabilities = data.get("server_capabilities") |
| 94 | + capability = capabilities.get("activity_heartbeat_resume") if isinstance(capabilities, Mapping) else None |
| 95 | + if (not isinstance(task, dict) or not isinstance(capability, Mapping) |
| 96 | + or capability.get("supported") is not True or type(capability.get("version")) is not int |
| 97 | + or capability["version"] != 1): |
| 98 | + return task |
| 99 | + if "last_heartbeat_details" not in task and "last_heartbeat_at" not in task: |
| 100 | + return task |
| 101 | + identity = _identity(task) |
| 102 | + if task["namespace"] != namespace or task.get("lease_owner") != worker_id: |
| 103 | + raise _invalid() |
| 104 | + if task_queue is not None and "task_queue" in task and task["task_queue"] != task_queue: |
| 105 | + raise _invalid() |
| 106 | + if "last_heartbeat_details" not in task or "last_heartbeat_at" not in task: |
| 107 | + raise _invalid() |
| 108 | + details, at = task["last_heartbeat_details"], task["last_heartbeat_at"] |
| 109 | + if details is None and at is None: |
| 110 | + checkpoint = None |
| 111 | + elif details is not None and task["attempt_number"] > 1: |
| 112 | + checkpoint = ActivityHeartbeatCheckpoint(details, at) |
| 113 | + else: |
| 114 | + raise _invalid() |
| 115 | + return _ActivityTask(task, _HeartbeatResume(True, checkpoint, identity)) |
| 116 | + |
| 117 | + |
| 118 | +def claim_resume( |
| 119 | + task: Mapping[str, Any], namespace: str | None, task_id: str, attempt_id: str, attempt: int, worker_id: str, |
| 120 | + activity_type: str, task_queue: str, |
| 121 | +) -> _HeartbeatResume: |
| 122 | + if not isinstance(task, _ActivityTask): |
| 123 | + return _HeartbeatResume() |
| 124 | + resume = task.heartbeat_resume |
| 125 | + if (resume.identity != _identity(task) or task.get("namespace") != namespace |
| 126 | + or task.get("task_id") != task_id or task.get("activity_attempt_id") != attempt_id |
| 127 | + or type(attempt) is not int or task.get("attempt_number") != attempt |
| 128 | + or task.get("lease_owner") != worker_id or task.get("activity_type") != activity_type |
| 129 | + or "task_queue" in task and task["task_queue"] != task_queue): |
| 130 | + raise _invalid() |
| 131 | + return resume |
0 commit comments