Add TaskStage API support with models, resource, tests and examples

examples/task_stage_example.py

@@ -0,0 +1,60 @@
+"""
+Example usage of TaskStages API
+
+Demonstrates:
+- Read a task stage
+- List task stages for a run
+- Override a task stage
+"""
+
+import os
+import sys
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
+
+from pytfe import TFEClient, TFEConfig
+
+
+def main():
+    client = TFEClient(TFEConfig.from_env())
+
+    task_stage_id = os.getenv("TFE_TASK_STAGE_ID")
+    run_id = os.getenv("TFE_RUN_ID")
+
+    if not task_stage_id or not run_id:
+        print("Please set TFE_TASK_STAGE_ID and TFE_RUN_ID")
+        return
+
+    print("=== TaskStages Example ===")
+
+    # READ
+    print("\nReading task stage...")
+    try:
+        stage = client.task_stages.read(task_stage_id)
+        print(f"ID:     {stage.id}")
+        print(f"Stage:  {stage.stage}")
+        print(f"Status: {stage.status}")
+        print(f"Run:    {stage.run.id if stage.run else None}")
+    except Exception as e:
+        print(f"Read failed: {e}")
+
+    # LIST
+    print("\nListing task stages...")
+    try:
+        stages = list(client.task_stages.list(run_id))
+        for s in stages:
+            print(f"{s.id} - {s.status}")
+    except Exception as e:
+        print(f"List failed: {e}")
+
+    # OVERRIDE
+    print("\nOverriding task stage...")
+    try:
+        client.task_stages.override(task_stage_id, comment="Approved")
+        print("Override successful")
+    except Exception as e:
+        print(f"Override failed: {e}")
+
+
+if __name__ == "__main__":
+    main()

src/pytfe/client.py

@@ -44,6 +44,7 @@
 from .resources.stack_configuration import StackConfigurations
 from .resources.state_version_outputs import StateVersionOutputs
 from .resources.state_versions import StateVersions
+from .resources.task_stage import TaskStages
 from .resources.task_result import TaskResults
 from .resources.team import Teams
 from .resources.team_project_access import TeamProjectAccesses
@@ -119,6 +120,7 @@ def __init__(self, config: TFEConfig | None = None):
         self.run_task_integrations = RunTaskIntegrations(self._transport)
         self.run_triggers = RunTriggers(self._transport)
         self.runs = Runs(self._transport)
+        self.task_stages = TaskStages(self._transport)
         self.query_runs = QueryRuns(self._transport)
         self.run_events = RunEvents(self._transport)
         self.comments = Comments(self._transport)

src/pytfe/errors.py

@@ -490,7 +490,7 @@ def __init__(self, message: str = "must provide at least one policy"):
 
 # Policy Evaluation errors
 class InvalidTaskStageIDError(InvalidValues):
-    """Raised when an invalid task stage ID is provided."""
+    """Raised when a task stage ID is invalid."""
 
     def __init__(self, message: str = "invalid value for task stage ID"):
         super().__init__(message)

src/pytfe/models/__init__.py

@@ -367,6 +367,8 @@
     StateVersionOutputsListOptions,
 )
 
+# ── Task Stage & Task Result ─────────────────────────────────────────────────
+from .task_result import TaskResult
 # ── Task Result ───────────────────────────────────────────────────────────────
 from .task_result import (
     TaskEnforcementLevel as TaskResultEnforcementLevel,
@@ -713,6 +715,9 @@
     "RunEventList",
     "RunEventListOptions",
     "RunEventReadOptions",
+    # Task Stage & Task Result
+    "TaskStage",
+    "TaskResult",
     # Comments
     "Comment",
     "CommentCreateOptions",
@@ -826,6 +831,13 @@
 RegistryProvider.model_rebuild()
 RegistryProviderVersion.model_rebuild()
 RegistryProviderPlatform.model_rebuild()
+TaskStage.model_rebuild(
+    raise_errors=False,
+    _types_namespace={
+        "Run": Run,
+        "TaskResult": TaskResult,
+        "PolicyEvaluation": PolicyEvaluation,
+
 
 # Rebuild TaskResult to resolve Run, Workspace, PolicyEvaluation, TaskStage refs
 TaskResult.model_rebuild(
@@ -835,5 +847,6 @@
         "Run": Run,
         "TaskStage": TaskStage,
         "Workspace": Workspace,
+
     },
 )

src/pytfe/models/task_result.py

@@ -10,6 +10,7 @@
 from pydantic import BaseModel, ConfigDict, Field
 
 if TYPE_CHECKING:
+    from pytfe.models.task_stage import TaskStage
     # Imported only for type checking to avoid circular imports.
     from pytfe.models.policy_evaluation import PolicyEvaluation
     from pytfe.models.run import Run
@@ -32,8 +33,8 @@ class TaskEnforcementLevel(str, Enum):
 
 
 class TaskResultStatusTimestamps(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
     model_config = ConfigDict(populate_by_name=True)
-
     errored_at: datetime | None = Field(None, alias="errored-at")
     running_at: datetime | None = Field(None, alias="running-at")
     canceled_at: datetime | None = Field(None, alias="canceled-at")
@@ -42,6 +43,9 @@ class TaskResultStatusTimestamps(BaseModel):
 
 
 class TaskResult(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
+
+    id: str
     model_config = ConfigDict(populate_by_name=True)
 
     id: str
@@ -50,6 +54,8 @@ class TaskResult(BaseModel):
     message: str | None = Field(None, alias="message")
 
     status_timestamps: TaskResultStatusTimestamps | None = Field(
+        None, alias="status-timestamps"
+
         None,
         alias="status-timestamps",
     )
@@ -64,6 +70,15 @@ class TaskResult(BaseModel):
     task_url: str | None = Field(None, alias="task-url")
 
     workspace_task_id: str | None = Field(None, alias="workspace-task-id")
+    workspace_task_enforcement_level: TaskEnforcementLevel | None = Field(
+        None, alias="workspace-task-enforcement-level"
+    )
+
+    agent_pool_id: str | None = Field(None, alias="agent-pool-id")
+    task_stage: TaskStage | None = Field(None, alias="task-stage")
+
+    @classmethod
+    def model_validate(cls, *args: Any, **kwargs: Any) -> TaskResult:
 
     workspace_task_enforcement_level: TaskEnforcementLevel | None = Field(
         None,
@@ -93,7 +108,14 @@ def model_validate(cls, *args: Any, **kwargs: Any) -> TaskResult:
 
 
 def _rebuild_task_result_model() -> None:
-    # Resolve all forward references once all modules are loaded.
+    try:
+        from pytfe.models.task_stage import TaskStage
+
+        TaskResult.model_rebuild(
+            raise_errors=False,
+            _types_namespace={"TaskStage": TaskStage},
+        )
+    except Exception:
     try:
         from pytfe.models.policy_evaluation import PolicyEvaluation
         from pytfe.models.run import Run
@@ -110,7 +132,6 @@ def _rebuild_task_result_model() -> None:
             },
         )
     except Exception:
-        # One or more models not yet importable during partial init; safe to skip.
         pass
 
 

src/pytfe/models/task_stage.py

@@ -3,23 +3,98 @@
 
 from __future__ import annotations
 
-from pydantic import BaseModel, ConfigDict
+from datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from pytfe.models.policy_evaluation import PolicyEvaluation
+from pytfe.models.task_result import TaskResult
+
+if TYPE_CHECKING:
+    from pytfe.models.run import Run
+
+
+class Stage(str, Enum):
+    pre_plan = "pre_plan"
+    post_plan = "post_plan"
+    pre_apply = "pre_apply"
+    post_apply = "post_apply"
+
+
+class TaskStageStatus(str, Enum):
+    pending = "pending"
+    running = "running"
+    passed = "passed"
+    failed = "failed"
+    awaiting_override = "awaiting_override"
+    canceled = "canceled"
+    errored = "errored"
+    unreachable = "unreachable"
+
+
+class TaskStageStatusTimestamps(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
+
+    errored_at: datetime | None = Field(None, alias="errored-at")
+    running_at: datetime | None = Field(None, alias="running-at")
+    canceled_at: datetime | None = Field(None, alias="canceled-at")
+    failed_at: datetime | None = Field(None, alias="failed-at")
+    passed_at: datetime | None = Field(None, alias="passed-at")
+
+
+class Permissions(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
+
+    can_override_policy: bool | None = Field(None, alias="can-override-policy")
+    can_override_tasks: bool | None = Field(None, alias="can-override-tasks")
+    can_override: bool | None = Field(None, alias="can-override")
+
+
+class Actions(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
+
+    is_overridable: bool | None = Field(None, alias="is-overridable")
 
 
-# TaskStage represents a HCP Terraform or Terraform Enterprise run's stage where run tasks can occur
 class TaskStage(BaseModel):
     model_config = ConfigDict(populate_by_name=True)
 
     id: str
-    # stage: Stage = Field(..., alias="stage")
-    # status: TaskStageStatus = Field(..., alias="status")
-    # status_timestamps: TaskStageStatusTimestamps = Field(..., alias="status-timestamps")
-    # created_at: datetime = Field(..., alias="created-at")
-    # updated_at: datetime = Field(..., alias="updated-at")
-    # permissions: Permissions = Field(..., alias="permissions")
-    # actions: Actions = Field(..., alias="actions")
-
-    # # Relations
-    # run: Run = Field(..., alias="run")
-    # task_results: list[TaskResult] = Field(..., alias="task-results")
-    # policy_evaluations: list[PolicyEvaluation] = Field(..., alias="policy-evaluations")
+
+    stage: Stage | None = Field(None, alias="stage")
+    status: TaskStageStatus | None = Field(None, alias="status")
+    status_timestamps: TaskStageStatusTimestamps | None = Field(
+        None, alias="status-timestamps"
+    )
+    created_at: datetime | None = Field(None, alias="created-at")
+    updated_at: datetime | None = Field(None, alias="updated-at")
+    permissions: Permissions | None = Field(None, alias="permissions")
+    actions: Actions | None = Field(None, alias="actions")
+
+    # Relationships
+    run: Run | None = Field(None, alias="run")
+    task_results: list[TaskResult] | None = Field(None, alias="task-results")
+    policy_evaluations: list[PolicyEvaluation] | None = Field(
+        None, alias="policy-evaluations"
+    )
+
+
+def _rebuild_task_stage_model() -> None:
+    try:
+        from pytfe.models.run import Run
+
+        TaskStage.model_rebuild(
+            raise_errors=False,
+            _types_namespace={
+                "Run": Run,
+                "TaskResult": TaskResult,
+                "PolicyEvaluation": PolicyEvaluation,
+            },
+        )
+    except Exception:
+        pass
+
+
+_rebuild_task_stage_model()