parkhojeong commented on code in PR #66218:
URL: https://github.com/apache/airflow/pull/66218#discussion_r3181772584
##########
airflow-ctl/src/airflowctl/api/operations.py:
##########
@@ -920,3 +920,77 @@ def list_import_errors(self) ->
PluginImportErrorCollectionResponse | ServerResp
return
PluginImportErrorCollectionResponse.model_validate_json(self.response.content)
except ServerResponseError as e:
raise e
+
+
+class TaskOperations(BaseOperations):
+ """Task instance operations."""
+
+ def states_for_dag_run(
+ self, dag_id: str, dag_run_id: str, limit: int = 100
+ ) -> TaskInstanceCollectionResponse | ServerResponseError:
+ """
+ Get task instance states for a DAG run.
+
+ Args:
+ dag_id: The DAG ID.
+ dag_run_id: The DAG run ID.
+ limit: Maximum number of task instances to return.
+ """
+ try:
+ self.response = self.client.get(
+ f"/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances",
+ params={"limit": limit},
+ )
+ return
TaskInstanceCollectionResponse.model_validate_json(self.response.content)
+ except ServerResponseError as e:
+ raise e
+
+ def clear(
+ self,
+ dag_id: str,
+ dry_run: bool = False,
+ only_failed: bool = True,
+ only_running: bool = False,
+ reset_dag_runs: bool = True,
+ task_ids: list[str] | None = None,
+ dag_run_id: str | None = None,
+ include_upstream: bool = False,
+ include_downstream: bool = False,
+ include_future: bool = False,
+ include_past: bool = False,
+ ) -> ClearTaskInstanceCollectionResponse | ServerResponseError:
+ """
+ Clear task instances for a Dag run.
+
+ Args:
+ dag_id: The Dag ID.
+ dry_run: Perform a dry run (no actual changes).
+ only_failed: Only clear failed task instances.
+ only_running: Only clear running task instances.
+ reset_dag_runs: Reset Dag runs.
+ task_ids: Comma-separated list of task IDs to clear.
+ dag_run_id: The Dag run ID.
+ include_upstream: Include upstream tasks.
+ include_downstream: Include downstream tasks.
+ include_future: Include future Dag runs.
+ include_past: Include past Dag runs.
+ """
+ try:
+ body = {
+ "dry_run": dry_run,
+ "only_failed": only_failed,
+ "only_running": only_running,
+ "reset_dag_runs": reset_dag_runs,
+ "task_ids": task_ids,
+ "dag_run_id": dag_run_id,
+ "include_upstream": include_upstream,
+ "include_downstream": include_downstream,
+ "include_future": include_future,
+ "include_past": include_past,
+ }
+ self.response =
self.client.post(f"dags/{dag_id}/clearTaskInstances", json=body)
Review Comment:
```suggestion
self.response = self.client.post(
f"dags/{dag_id}/clearTaskInstances",
json=body.model_dump(mode="json", exclude_none=True,
exclude_unset=True),
)
```
and you can use it as body
##########
airflow-ctl/src/airflowctl/api/operations.py:
##########
@@ -920,3 +920,77 @@ def list_import_errors(self) ->
PluginImportErrorCollectionResponse | ServerResp
return
PluginImportErrorCollectionResponse.model_validate_json(self.response.content)
except ServerResponseError as e:
raise e
+
+
+class TaskOperations(BaseOperations):
+ """Task instance operations."""
+
+ def states_for_dag_run(
+ self, dag_id: str, dag_run_id: str, limit: int = 100
+ ) -> TaskInstanceCollectionResponse | ServerResponseError:
+ """
+ Get task instance states for a DAG run.
+
+ Args:
+ dag_id: The DAG ID.
+ dag_run_id: The DAG run ID.
+ limit: Maximum number of task instances to return.
+ """
+ try:
+ self.response = self.client.get(
+ f"/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances",
+ params={"limit": limit},
+ )
+ return
TaskInstanceCollectionResponse.model_validate_json(self.response.content)
+ except ServerResponseError as e:
+ raise e
+
+ def clear(
+ self,
+ dag_id: str,
+ dry_run: bool = False,
+ only_failed: bool = True,
+ only_running: bool = False,
+ reset_dag_runs: bool = True,
+ task_ids: list[str] | None = None,
+ dag_run_id: str | None = None,
+ include_upstream: bool = False,
+ include_downstream: bool = False,
+ include_future: bool = False,
+ include_past: bool = False,
+ ) -> ClearTaskInstanceCollectionResponse | ServerResponseError:
Review Comment:
```suggestion
def clear(
self, dag_id: str, body: ClearTaskInstancesBody
) -> ClearTaskInstanceCollectionResponse | ServerResponseError:
```
you can use ClearTaskInstancesBody
##########
airflow-ctl/src/airflowctl/api/operations.py:
##########
@@ -920,3 +920,77 @@ def list_import_errors(self) ->
PluginImportErrorCollectionResponse | ServerResp
return
PluginImportErrorCollectionResponse.model_validate_json(self.response.content)
except ServerResponseError as e:
raise e
+
+
+class TaskOperations(BaseOperations):
+ """Task instance operations."""
+
+ def states_for_dag_run(
+ self, dag_id: str, dag_run_id: str, limit: int = 100
+ ) -> TaskInstanceCollectionResponse | ServerResponseError:
+ """
+ Get task instance states for a DAG run.
+
+ Args:
+ dag_id: The DAG ID.
+ dag_run_id: The DAG run ID.
+ limit: Maximum number of task instances to return.
+ """
+ try:
+ self.response = self.client.get(
+ f"/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances",
+ params={"limit": limit},
+ )
+ return
TaskInstanceCollectionResponse.model_validate_json(self.response.content)
+ except ServerResponseError as e:
+ raise e
+
+ def clear(
+ self,
+ dag_id: str,
+ dry_run: bool = False,
+ only_failed: bool = True,
+ only_running: bool = False,
+ reset_dag_runs: bool = True,
+ task_ids: list[str] | None = None,
+ dag_run_id: str | None = None,
+ include_upstream: bool = False,
+ include_downstream: bool = False,
+ include_future: bool = False,
+ include_past: bool = False,
+ ) -> ClearTaskInstanceCollectionResponse | ServerResponseError:
Review Comment:
```suggestion
self, dag_id: str, body: ClearTaskInstancesBody
) -> ClearTaskInstanceCollectionResponse | ServerResponseError:
```
you can use ClearTaskInstancesBody
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
