Lee-W commented on code in PR #69397:
URL: https://github.com/apache/airflow/pull/69397#discussion_r3543317858
##########
airflow-ctl/src/airflowctl/api/operations.py:
##########
@@ -750,6 +752,45 @@ def list(self) -> ProviderCollectionResponse |
ServerResponseError:
return super().execute_list(path="providers",
data_model=ProviderCollectionResponse)
+class TaskInstancesOperations(BaseOperations):
+ """Task instance operations."""
+
+ def get(
+ self, dag_id: str, dag_run_id: str, task_id: str, map_index: int |
None = None
+ ) -> TaskInstanceResponse | ServerResponseError:
+ """Get a task instance for a Dag run."""
+ path = f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}"
Review Comment:
nit: Do you think we can extract this endpoint-building logic as a utility
function?
##########
airflow-ctl/tests/airflow_ctl/api/test_operations.py:
##########
@@ -1489,6 +1493,84 @@ def handle_request(request: httpx.Request) ->
httpx.Response:
assert response == self.provider_collection_response
+class TestTaskInstancesOperations:
+ task_instance_response = TaskInstanceResponse(
+ id=uuid.UUID("4d828a62-a417-4936-a7a6-2b3fabacecab"),
+ task_id="task_id",
+ dag_id="dag_id",
+ dag_run_id="dag_run_id",
+ map_index=-1,
+ run_after=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ state=TaskInstanceState.SUCCESS,
+ try_number=1,
+ max_tries=0,
+ task_display_name="task_id",
+ dag_display_name="dag_id",
+ pool="default_pool",
+ pool_slots=1,
+ executor_config="{}",
+ )
+ task_dependency_collection_response = TaskDependencyCollectionResponse(
+ dependencies=[TaskDependencyResponse(name="Trigger Rule",
reason="upstream tasks not done")],
+ )
+
+ @pytest.mark.parametrize("map_index", [None, -1])
+ def test_get(self, map_index):
+ def handle_request(request: httpx.Request) -> httpx.Response:
Review Comment:
IIRC, this practice is used in many places. But I'm kinda wondering whether
we can extract this definition? perfectly fine to skip it if it doesn't really
help
##########
airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py:
##########
@@ -0,0 +1,235 @@
+# 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.
+from __future__ import annotations
+
+import datetime
+import uuid
+from unittest import mock
+
+import httpx
+import pytest
+
+from airflowctl.api.datamodels.generated import (
+ TaskDependencyCollectionResponse,
+ TaskDependencyResponse,
+ TaskInstanceResponse,
+ TaskInstanceState,
+)
+from airflowctl.api.operations import ServerResponseError
+from airflowctl.ctl import cli_parser
+from airflowctl.ctl.commands import task_command
+
+
+def _server_error(status_code: int) -> ServerResponseError:
Review Comment:
```suggestion
def _make_server_error(status_code: int) -> ServerResponseError:
```
nit
##########
airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py:
##########
@@ -0,0 +1,235 @@
+# 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.
+from __future__ import annotations
+
+import datetime
+import uuid
+from unittest import mock
+
+import httpx
+import pytest
+
+from airflowctl.api.datamodels.generated import (
+ TaskDependencyCollectionResponse,
+ TaskDependencyResponse,
+ TaskInstanceResponse,
+ TaskInstanceState,
+)
+from airflowctl.api.operations import ServerResponseError
+from airflowctl.ctl import cli_parser
+from airflowctl.ctl.commands import task_command
+
+
+def _server_error(status_code: int) -> ServerResponseError:
+ request = httpx.Request("GET",
"http://testserver/api/v2/dags/test_dag/dagRuns/test_run";)
+ response = httpx.Response(status_code, request=request, json={"detail":
"boom"})
+ return ServerResponseError(message="boom", request=request,
response=response)
+
+
+class TestFailedDeps:
+ parser = cli_parser.get_parser()
+ dag_id = "test_dag"
+ run_id = "test_run"
+ task_id = "test_task"
+ logical_date = datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc)
+
+ def _make_task_instance(self, state: TaskInstanceState | None) ->
TaskInstanceResponse:
+ return TaskInstanceResponse(
+ id=uuid.uuid4(),
+ task_id=self.task_id,
+ dag_id=self.dag_id,
+ dag_run_id=self.run_id,
+ map_index=-1,
+ run_after=self.logical_date,
+ state=state,
+ try_number=1,
+ max_tries=0,
+ task_display_name=self.task_id,
+ dag_display_name=self.dag_id,
+ pool="default_pool",
+ pool_slots=1,
+ executor_config="{}",
+ )
+
+ def _make_api_client(
+ self,
+ dependencies: list[TaskDependencyResponse] | None = None,
+ state: TaskInstanceState | None = None,
+ ) -> mock.MagicMock:
+ api_client = mock.MagicMock()
+ api_client.dag_runs.get.return_value =
mock.MagicMock(dag_run_id=self.run_id)
+ api_client.task_instances.get_dependencies.return_value =
TaskDependencyCollectionResponse(
+ dependencies=dependencies or [],
+ )
+ api_client.task_instances.get.return_value =
self._make_task_instance(state=state)
+ return api_client
+
+ def test_failed_deps_not_met(self, capsys):
+ api_client = self._make_api_client(
+ dependencies=[
+ TaskDependencyResponse(
+ name="Dagrun Running", reason="Task instance's dagrun was
not in the 'running' state"
+ ),
+ TaskDependencyResponse(
+ name="Trigger Rule", reason="requires all upstream tasks
to have succeeded"
+ ),
+ ]
+ )
+
+ task_command.failed_deps(
+ self.parser.parse_args(["tasks", "failed-deps", self.dag_id,
self.task_id, self.run_id]),
+ api_client=api_client,
+ )
+
+ api_client.task_instances.get_dependencies.assert_called_once_with(
+ dag_id=self.dag_id,
+ dag_run_id=self.run_id,
+ task_id=self.task_id,
+ map_index=-1,
+ suppress_error_log=True,
+ )
+ api_client.task_instances.get.assert_not_called()
+ out = capsys.readouterr().out
+ assert "Task instance dependencies not met:" in out
+ assert "Dagrun Running: Task instance's dagrun was not in the
'running' state" in out
+ assert "Trigger Rule: requires all upstream tasks to have succeeded"
in out
+
Review Comment:
If the error message is not too long and hard to read, let's use the full
message
##########
airflow-ctl/src/airflowctl/ctl/utils/dag_run.py:
##########
@@ -0,0 +1,63 @@
+# 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.
+
+from __future__ import annotations
+
+import datetime
+from typing import TYPE_CHECKING
+
+from airflowctl.api.client import ServerResponseError
+
+if TYPE_CHECKING:
+ from airflowctl.api.client import Client
+ from airflowctl.api.datamodels.generated import DAGRunResponse
+
+
+def _parse_logical_date(value: str) -> datetime.datetime | None:
Review Comment:
I guess we'll need to rebase once the other PR is merged? look like the same
function as the other PR
##########
airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py:
##########
@@ -0,0 +1,235 @@
+# 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.
+from __future__ import annotations
+
+import datetime
+import uuid
+from unittest import mock
+
+import httpx
+import pytest
+
+from airflowctl.api.datamodels.generated import (
+ TaskDependencyCollectionResponse,
+ TaskDependencyResponse,
+ TaskInstanceResponse,
+ TaskInstanceState,
+)
+from airflowctl.api.operations import ServerResponseError
+from airflowctl.ctl import cli_parser
+from airflowctl.ctl.commands import task_command
+
+
+def _server_error(status_code: int) -> ServerResponseError:
+ request = httpx.Request("GET",
"http://testserver/api/v2/dags/test_dag/dagRuns/test_run";)
+ response = httpx.Response(status_code, request=request, json={"detail":
"boom"})
+ return ServerResponseError(message="boom", request=request,
response=response)
+
+
+class TestFailedDeps:
+ parser = cli_parser.get_parser()
+ dag_id = "test_dag"
+ run_id = "test_run"
+ task_id = "test_task"
+ logical_date = datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc)
+
+ def _make_task_instance(self, state: TaskInstanceState | None) ->
TaskInstanceResponse:
+ return TaskInstanceResponse(
+ id=uuid.uuid4(),
+ task_id=self.task_id,
+ dag_id=self.dag_id,
+ dag_run_id=self.run_id,
+ map_index=-1,
+ run_after=self.logical_date,
+ state=state,
+ try_number=1,
+ max_tries=0,
+ task_display_name=self.task_id,
+ dag_display_name=self.dag_id,
+ pool="default_pool",
+ pool_slots=1,
+ executor_config="{}",
+ )
+
+ def _make_api_client(
+ self,
+ dependencies: list[TaskDependencyResponse] | None = None,
+ state: TaskInstanceState | None = None,
+ ) -> mock.MagicMock:
+ api_client = mock.MagicMock()
+ api_client.dag_runs.get.return_value =
mock.MagicMock(dag_run_id=self.run_id)
+ api_client.task_instances.get_dependencies.return_value =
TaskDependencyCollectionResponse(
+ dependencies=dependencies or [],
+ )
+ api_client.task_instances.get.return_value =
self._make_task_instance(state=state)
+ return api_client
+
+ def test_failed_deps_not_met(self, capsys):
+ api_client = self._make_api_client(
+ dependencies=[
+ TaskDependencyResponse(
+ name="Dagrun Running", reason="Task instance's dagrun was
not in the 'running' state"
+ ),
+ TaskDependencyResponse(
+ name="Trigger Rule", reason="requires all upstream tasks
to have succeeded"
+ ),
+ ]
+ )
+
+ task_command.failed_deps(
+ self.parser.parse_args(["tasks", "failed-deps", self.dag_id,
self.task_id, self.run_id]),
+ api_client=api_client,
+ )
+
+ api_client.task_instances.get_dependencies.assert_called_once_with(
+ dag_id=self.dag_id,
+ dag_run_id=self.run_id,
+ task_id=self.task_id,
+ map_index=-1,
+ suppress_error_log=True,
+ )
+ api_client.task_instances.get.assert_not_called()
+ out = capsys.readouterr().out
+ assert "Task instance dependencies not met:" in out
+ assert "Dagrun Running: Task instance's dagrun was not in the
'running' state" in out
+ assert "Trigger Rule: requires all upstream tasks to have succeeded"
in out
+
Review Comment:
same can be applied to elsewhere
--
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]
