Type Stage 1 status snapshots

Author: anth-volk

policyengine_us_data/build_datasets/__init__.py

@@ -43,7 +43,9 @@
 from .status import Stage1ErrorRecord, Stage1StatusEvent
 from .status_store import (
     Stage1StatusRecorder,
+    Stage1StatusReadError,
     Stage1StatusSnapshot,
+    Stage1StoredStatusEvent,
     empty_stage_1_status_snapshot,
     read_stage_1_status_snapshot,
 )
@@ -69,9 +71,11 @@
     "Stage1Coordinator",
     "Stage1ErrorRecord",
     "Stage1StatusRecorder",
+    "Stage1StatusReadError",
     "Stage1StatusEvent",
     "Stage1StatusSink",
     "Stage1StatusSnapshot",
+    "Stage1StoredStatusEvent",
     "Stage1SubstepRunner",
     "SubprocessLogCapture",
     "TargetDatabaseSchemaSummaryWriter",

policyengine_us_data/build_datasets/results.py

@@ -2,9 +2,9 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 from .status import Stage1ErrorRecord, Stage1SubstepStatus
 
@@ -27,6 +27,24 @@ class DatasetCommandResult:
     error: Stage1ErrorRecord | None = None
     metadata: Mapping[str, Any] = field(default_factory=dict)
 
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> "DatasetCommandResult":
+        """Build a command result from a JSON-compatible payload."""
+
+        error = _error_record_from_payload(data.get("error"))
+        return cls(
+            command_name=str(data["command_name"]),
+            argv=_string_tuple(data.get("argv", ())),
+            status=_command_execution_status(data["status"]),
+            returncode=_optional_int(data.get("returncode")),
+            started_at=str(data["started_at"]),
+            completed_at=str(data["completed_at"]),
+            duration_s=float(data["duration_s"]),
+            combined_output_tail=_string_tuple(data.get("combined_output_tail", ())),
+            error=error,
+            metadata=_metadata_mapping(data.get("metadata", {})),
+        )
+
     def to_dict(self) -> dict[str, Any]:
         """Return a JSON-compatible command result payload."""
 
@@ -60,6 +78,32 @@ class DatasetSubstepResult:
     error: Stage1ErrorRecord | None = None
     metadata: Mapping[str, Any] = field(default_factory=dict)
 
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> "DatasetSubstepResult":
+        """Build a substep result from a JSON-compatible payload."""
+
+        command_results = data.get("command_results", ())
+        if not isinstance(command_results, Sequence) or isinstance(
+            command_results, str
+        ):
+            raise TypeError("command_results must be a sequence")
+        return cls(
+            substep_id=str(data["substep_id"]),
+            title=str(data["title"]),
+            status=_stage_1_substep_status(data["status"]),
+            started_at=_optional_str(data.get("started_at")),
+            completed_at=str(data["completed_at"]),
+            duration_s=_optional_float(data.get("duration_s")),
+            command_names=_string_tuple(data.get("command_names", ())),
+            command_results=tuple(
+                DatasetCommandResult.from_dict(_mapping_payload(result))
+                for result in command_results
+            ),
+            artifact_paths=_string_tuple(data.get("artifact_paths", ())),
+            error=_error_record_from_payload(data.get("error")),
+            metadata=_metadata_mapping(data.get("metadata", {})),
+        )
+
     def to_dict(self) -> dict[str, Any]:
         """Return a JSON-compatible substep result payload."""
 
@@ -78,6 +122,54 @@ def to_dict(self) -> dict[str, Any]:
         }
 
 
+def _command_execution_status(value: Any) -> CommandExecutionStatus:
+    if value in ("completed", "failed"):
+        return cast(CommandExecutionStatus, value)
+    raise ValueError(f"Invalid command execution status: {value!r}")
+
+
+def _stage_1_substep_status(value: Any) -> Stage1SubstepStatus:
+    if value in ("started", "completed", "skipped", "failed"):
+        return cast(Stage1SubstepStatus, value)
+    raise ValueError(f"Invalid Stage 1 substep status: {value!r}")
+
+
+def _string_tuple(value: Any) -> tuple[str, ...]:
+    if not isinstance(value, Sequence) or isinstance(value, str):
+        raise TypeError("Expected a sequence")
+    return tuple(str(item) for item in value)
+
+
+def _mapping_payload(value: Any) -> Mapping[str, Any]:
+    if not isinstance(value, Mapping):
+        raise TypeError("Expected a mapping")
+    return value
+
+
+def _metadata_mapping(value: Any) -> Mapping[str, Any]:
+    if not isinstance(value, Mapping):
+        raise TypeError("metadata must be a mapping")
+    return dict(value)
+
+
+def _error_record_from_payload(value: Any) -> Stage1ErrorRecord | None:
+    if value is None:
+        return None
+    return Stage1ErrorRecord.from_dict(_mapping_payload(value))
+
+
+def _optional_str(value: Any) -> str | None:
+    return None if value is None else str(value)
+
+
+def _optional_int(value: Any) -> int | None:
+    return None if value is None else int(value)
+
+
+def _optional_float(value: Any) -> float | None:
+    return None if value is None else float(value)
+
+
 __all__ = [
     "CommandExecutionStatus",
     "DatasetCommandResult",

policyengine_us_data/build_datasets/status.py

@@ -5,7 +5,7 @@
 from collections.abc import Mapping
 from dataclasses import dataclass, field, replace
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 if TYPE_CHECKING:
     from modal_app.step_manifests.errors import PipelineErrorRecord
@@ -37,6 +37,19 @@ class Stage1StatusEvent:
     command_name: str | None = None
     metadata: Mapping[str, Any] = field(default_factory=dict)
 
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> "Stage1StatusEvent":
+        """Build a status event from a JSON-compatible payload."""
+
+        return cls(
+            substep_id=str(data["substep_id"]),
+            status=_stage_1_substep_status(data["status"]),
+            created_at=str(data["created_at"]),
+            message=_optional_str(data.get("message")),
+            command_name=_optional_str(data.get("command_name")),
+            metadata=_metadata_mapping(data.get("metadata", {})),
+        )
+
     def to_dict(self) -> dict[str, Any]:
         """Return a JSON-compatible status event payload."""
 
@@ -67,6 +80,20 @@ class Stage1ErrorRecord:
     created_at: str = field(default_factory=utc_timestamp)
     metadata: Mapping[str, Any] = field(default_factory=dict)
 
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> "Stage1ErrorRecord":
+        """Build an error record from a JSON-compatible payload."""
+
+        return cls(
+            substep_id=_optional_str(data.get("substep_id")),
+            command_name=_optional_str(data.get("command_name")),
+            error_type=str(data["error_type"]),
+            message=str(data["message"]),
+            returncode=_optional_int(data.get("returncode")),
+            created_at=str(data.get("created_at") or utc_timestamp()),
+            metadata=_metadata_mapping(data.get("metadata", {})),
+        )
+
     @classmethod
     def from_exception(
         cls,
@@ -159,6 +186,26 @@ def _pipeline_traceback_text(error: Stage1ErrorRecord) -> str:
     return "\n\n".join(parts)
 
 
+def _stage_1_substep_status(value: Any) -> Stage1SubstepStatus:
+    if value in ("started", "completed", "skipped", "failed"):
+        return cast(Stage1SubstepStatus, value)
+    raise ValueError(f"Invalid Stage 1 substep status: {value!r}")
+
+
+def _optional_str(value: Any) -> str | None:
+    return None if value is None else str(value)
+
+
+def _optional_int(value: Any) -> int | None:
+    return None if value is None else int(value)
+
+
+def _metadata_mapping(value: Any) -> Mapping[str, Any]:
+    if not isinstance(value, Mapping):
+        raise TypeError("metadata must be a mapping")
+    return dict(value)
+
+
 __all__ = [
     "Stage1ErrorRecord",
     "Stage1StatusEvent",