parkhojeong commented on code in PR #66218:
URL: https://github.com/apache/airflow/pull/66218#discussion_r3192736414
##########
airflow-ctl/src/airflowctl/ctl/cli_config.py:
##########
@@ -944,6 +944,66 @@ def merge_commands(
),
)
+
+TASKS_COMMANDS = (
+ ActionCommand(
+ name="clear",
+ help="Clear task instances for a Dag run",
+ func=lazy_load_command("airflowctl.ctl.commands.task_command.clear"),
+ args=(
+ ARG_DAG_ID,
+ ARG_OUTPUT,
+ ARG_DAG_RUN_ID,
Review Comment:
<img width="534" height="259" alt="Image"
src="https://github.com/user-attachments/assets/b05e7b04-acd9-46b6-ba44-79023039e719";
/>
ARG_DAG_RUN_ID doesn't seem to be defined, so importing cli_config fails
with NameError.
##########
airflow-ctl/src/airflowctl/api/operations.py:
##########
@@ -920,3 +920,50 @@ 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,
+ body: ClearTaskInstancesBody,
+ ) -> TaskInstanceCollectionResponse | ServerResponseError:
+ """
+ Clear task instances for a Dag run.
+
+ Args:
+ dag_id: The Dag ID.
+ body: ClearTaskInstancesBody containing the clear parameters.
+ """
Review Comment:
Nit: I don’t think a comment is necessary here, since the code is already
clear enough. What do you think?
##########
airflow-ctl/src/airflowctl/api/operations.py:
##########
@@ -920,3 +920,50 @@ 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:
Review Comment:
<img width="680" height="313" alt="Image"
src="https://github.com/user-attachments/assets/da769f18-b6e7-45d1-9357-e88728b9d746";
/>
TaskInstanceCollectionResponse doesn't seems to be imported
##########
airflow-ctl/src/airflowctl/ctl/commands/task_command.py:
##########
@@ -0,0 +1,82 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Task commands for airflowctl."""
+from __future__ import annotations
+
+import sys
+
+import rich
+
+from airflowctl.api.client import (
+ NEW_API_CLIENT,
+ ClientKind,
+ ServerResponseError,
+ provide_api_client,
+)
+from airflowctl.api.datamodels.generated import (
+ ClearTaskInstancesBody,
+ ClearTaskInstanceCollectionResponse,
+)
+
+from airflowctl.ctl.console_formatting import AirflowConsole
+
+
+@provide_api_client(kind=ClientKind.CLI)
+def clear(args, api_client=NEW_API_CLIENT) -> None:
+ """Clear task instances for a Dag run."""
+ try:
+ task_ids = getattr(args, "task_ids", None)
+ if task_ids:
+ task_ids = task_ids.split(",")
+
+ body = ClearTaskInstancesBody(
+ dry_run=getattr(args, "dry_run", False) or False,
+ only_failed=getattr(args, "only_failed", True) or True,
+ only_running=getattr(args, "only_running", False) or False,
+ reset_dag_runs=getattr(args, "reset_dag_runs", True) or True,
+ task_ids=task_ids,
+ dag_run_id=getattr(args, "dag_run_id", None),
+ include_upstream=getattr(args, "upstream", False) or False,
+ include_downstream=getattr(args, "downstream", False) or False,
+ include_future=getattr(args, "include_future", False) or False,
+ include_past=getattr(args, "include_past", False) or False,
+ )
Review Comment:
```suggestion
body = ClearTaskInstancesBody(
dry_run=getattr(args, "dry_run", False),
only_failed=getattr(args, "only_failed", True),
only_running=getattr(args, "only_running", False),
reset_dag_runs=getattr(args, "reset_dag_runs", True),
task_ids=task_ids,
dag_run_id=getattr(args, "dag_run_id", None),
include_upstream=getattr(args, "upstream", False),
include_downstream=getattr(args, "downstream", False),
include_future=getattr(args, "include_future", False),
include_past=getattr(args, "include_past", False),
)
```
The trailing or False/or True expressions seem unnecessary here. Is there a
reason we need them?
##########
airflow-ctl/src/airflowctl/ctl/commands/task_command.py:
##########
@@ -0,0 +1,82 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Task commands for airflowctl."""
+from __future__ import annotations
+
+import sys
+
+import rich
+
+from airflowctl.api.client import (
+ NEW_API_CLIENT,
+ ClientKind,
+ ServerResponseError,
+ provide_api_client,
+)
+from airflowctl.api.datamodels.generated import (
+ ClearTaskInstancesBody,
+ ClearTaskInstanceCollectionResponse,
+)
+
+from airflowctl.ctl.console_formatting import AirflowConsole
+
+
+@provide_api_client(kind=ClientKind.CLI)
+def clear(args, api_client=NEW_API_CLIENT) -> None:
+ """Clear task instances for a Dag run."""
+ try:
+ task_ids = getattr(args, "task_ids", None)
+ if task_ids:
+ task_ids = task_ids.split(",")
+
+ body = ClearTaskInstancesBody(
+ dry_run=getattr(args, "dry_run", False) or False,
+ only_failed=getattr(args, "only_failed", True) or True,
+ only_running=getattr(args, "only_running", False) or False,
+ reset_dag_runs=getattr(args, "reset_dag_runs", True) or True,
+ task_ids=task_ids,
+ dag_run_id=getattr(args, "dag_run_id", None),
+ include_upstream=getattr(args, "upstream", False) or False,
+ include_downstream=getattr(args, "downstream", False) or False,
+ include_future=getattr(args, "include_future", False) or False,
+ include_past=getattr(args, "include_past", False) or False,
+ )
+ cleared = api_client.client.tasks.clear(dag_id=args.dag_id, body=body)
+ rich.print(
+ f"[green]Cleared {len(cleared.task_instances)} task instance(s)
for Dag {args.dag_id}[/green]"
+ )
+ AirflowConsole().print_as(
+ data=[t.model_dump() for t in cleared.task_instances],
+ output=args.output,
+ )
+ except ServerResponseError as e:
+ rich.print(f"[red]Error clearing task instances: {e}[/red]")
+ sys.exit(1)
+
+@provide_api_client(kind=ClientKind.CLI)
+def states_for_dag_run(args, api_client=NEW_API_CLIENT) -> None:
+ """Get task instance states for a Dag run."""
+ try:
+ tis = api_client.client.tasks.list(dag_id=args.dag_id,
dag_run_id=args.dag_run_id)
Review Comment:
```suggestion
tis = api_client.client.tasks.states_for_dag_run(dag_id=args.dag_id,
dag_run_id=args.dag_run_id)
```
`TaskOperations` defines `states_for_dag_run`, but this calls `tasks.list`,
which does not seem to exist. This would fail at runtime, and it also ignores
the `--limit` option added for this command.
##########
airflow-ctl/tests/airflow_ctl/api/test_operations.py:
##########
@@ -1825,6 +1825,89 @@ def test_list(self):
def handle_request(request: httpx.Request) -> httpx.Response:
assert request.url.path == ("/api/v2/plugins")
+
+
+class TestTaskOperations:
+ """Unit tests for TaskOperations."""
+
+ def test_states_for_dag_run(self):
+ """Test states_for_dag_run calls the correct endpoint with correct
params."""
+ import httpx
+ from airflowctl.api.datamodels.generated import
TaskInstanceCollectionResponse, TaskInstanceResponse
+
+ task_response = TaskInstanceCollectionResponse(
+ task_instances=[
+ TaskInstanceResponse(
+ task_id="task_1",
+ dag_id="example_dag",
+ dag_run_id="run_1",
+ logical_date="2025-01-01T00:00:00Z",
+ start_date="2025-01-01T00:01:00Z",
+ end_date="2025-01-01T00:05:00Z",
+ duration=240.0,
+ state="success",
+ try_number=1,
+ max_tries=2,
+ task_display_name="Task 1",
+ )
+ ],
+ total_entries=1,
+ )
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ assert request.url.path ==
"/dags/example_dag/dagRuns/run_1/taskInstances"
+ params = dict(request.url.params)
+ assert params.get("limit") == "10"
+ return httpx.Response(200,
json=json.loads(task_response.model_dump_json()))
+
+ from airflowctl.api.operations import make_api_client, TaskOperations
+ client = make_api_client(transport=httpx.MockTransport(handle_request))
+ ops = TaskOperations(client=client.client)
+ response = ops.states_for_dag_run("example_dag", "run_1", limit=10)
+ assert response.total_entries == 1
+ assert len(response.task_instances) == 1
+ assert response.task_instances[0].task_id == "task_1"
+
+ def test_clear(self):
+ """Test clear calls the correct endpoint with correct params."""
+ import httpx
+ from airflowctl.api.datamodels.generated import
ClearTaskInstancesBody, TaskInstanceCollectionResponse, TaskInstanceResponse
+
+ clear_response = TaskInstanceCollectionResponse(
+ task_instances=[
+ TaskInstanceResponse(
+ task_id="task_1",
+ dag_id="example_dag",
+ dag_run_id="run_1",
+ logical_date="2025-01-01T00:00:00Z",
+ start_date="2025-01-01T00:01:00Z",
+ end_date="2025-01-01T00:05:00Z",
+ duration=240.0,
+ state="cleared",
+ try_number=1,
+ max_tries=2,
+ task_display_name="Task 1",
+ )
+ ],
+ total_entries=1,
+ )
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ assert request.url.path == "/dags/example_dag/clearTaskInstances"
+ body = json.loads(request.content)
+ assert body["only_failed"] is True
+ return httpx.Response(200,
json=json.loads(clear_response.model_dump_json()))
+
+ from airflowctl.api.operations import make_api_client, TaskOperations
+ client = make_api_client(transport=httpx.MockTransport(handle_request))
+ ops = TaskOperations(client=client.client)
+ body = ClearTaskInstancesBody(only_failed=True)
+ response = ops.clear("example_dag", body=body)
+ assert response.total_entries == 1
+ assert len(response.task_instances) == 1
+ assert response.task_instances[0].state == "cleared"
+
+
return httpx.Response(200,
json=json.loads(self.plugin_collection_response.model_dump_json()))
Review Comment:
The tests can't run because of an indentation error.
I'm not sure why this return statement is here. Is there a reason we need
it?
##########
airflow-ctl/src/airflowctl/ctl/commands/task_command.py:
##########
@@ -0,0 +1,82 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Task commands for airflowctl."""
+from __future__ import annotations
+
+import sys
+
+import rich
+
+from airflowctl.api.client import (
+ NEW_API_CLIENT,
+ ClientKind,
+ ServerResponseError,
+ provide_api_client,
+)
+from airflowctl.api.datamodels.generated import (
+ ClearTaskInstancesBody,
+ ClearTaskInstanceCollectionResponse,
+)
+
+from airflowctl.ctl.console_formatting import AirflowConsole
+
+
+@provide_api_client(kind=ClientKind.CLI)
+def clear(args, api_client=NEW_API_CLIENT) -> None:
+ """Clear task instances for a Dag run."""
+ try:
+ task_ids = getattr(args, "task_ids", None)
+ if task_ids:
+ task_ids = task_ids.split(",")
+
+ body = ClearTaskInstancesBody(
+ dry_run=getattr(args, "dry_run", False) or False,
Review Comment:
```suggestion
body = ClearTaskInstancesBody(
dry_run=getattr(args, "dry_run", True),
```
Would it be better for the CLI to follow the API default (True) and only
send dry_run=False when the user explicitly disables dry-run mode?
This may also need a corresponding CLI option for explicitly performing the
non-dry-run clear.
For reference, `ClearTaskInstancesBody` defaults `dry_run` to `true` in
`airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py:203`:
```python
class ClearTaskInstancesBody(StrictBaseModel):
dry_run: bool = True
```
--
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]
