From 9ebd57a4c1a290979cb7c787dd9f552dfdfa92ce Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Fri, 24 Jul 2026 20:20:14 +0900 Subject: [PATCH] Add airflowctl tasks failed-deps command Signed-off-by: PoAn Yang --- .../src/airflow/cli/commands/task_command.py | 1 + .../cli/commands/test_command_deprecations.py | 1 + .../test_airflowctl_commands.py | 4 + airflow-ctl/docs/images/command_hashes.txt | 4 +- .../docs/images/output_taskinstances.svg | 74 +++-- airflow-ctl/docs/images/output_tasks.svg | 74 +++-- airflow-ctl/src/airflowctl/api/operations.py | 48 +++ airflow-ctl/src/airflowctl/ctl/cli_config.py | 28 ++ .../airflowctl/ctl/commands/task_command.py | 52 +++ .../src/airflowctl/ctl/help_texts.yaml | 2 + .../tests/airflow_ctl/api/test_operations.py | 60 +++- .../ctl/commands/test_task_command.py | 314 ++++++++++++++++++ .../tests/airflow_ctl/ctl/test_cli_config.py | 2 + 13 files changed, 590 insertions(+), 74 deletions(-) diff --git a/airflow-core/src/airflow/cli/commands/task_command.py b/airflow-core/src/airflow/cli/commands/task_command.py index b8cae95ae8087..5c37710e04694 100644 --- a/airflow-core/src/airflow/cli/commands/task_command.py +++ b/airflow-core/src/airflow/cli/commands/task_command.py @@ -262,6 +262,7 @@ class TaskCommandMarker: """Marker for listener hooks, to properly detect from which component they are called.""" +@deprecated_for_airflowctl("airflowctl tasks failed-deps") @cli_utils.action_cli(check_db=False) @providers_configuration_loaded def task_failed_deps(args) -> None: diff --git a/airflow-core/tests/unit/cli/commands/test_command_deprecations.py b/airflow-core/tests/unit/cli/commands/test_command_deprecations.py index 5e9d0613acd07..242d9d8ad723f 100644 --- a/airflow-core/tests/unit/cli/commands/test_command_deprecations.py +++ b/airflow-core/tests/unit/cli/commands/test_command_deprecations.py @@ -79,6 +79,7 @@ (config_command.show_config, "airflowctl config list"), (task_command.task_states_for_dag_run, "airflowctl tasks states-for-dag-run"), (task_command.task_clear, "airflowctl tasks clear"), + (task_command.task_failed_deps, "airflowctl tasks failed-deps"), ] diff --git a/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py b/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py index d269c5bb03282..14a500c4d5bf7 100644 --- a/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py +++ b/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py @@ -97,10 +97,14 @@ def date_param(): # Dag Run commands "dagrun list --dag-id example_bash_operator --state success --limit=1", # Tasks commands + 'tasks failed-deps example_bash_operator runme_0 "manual__{date_param}"', + 'tasks failed-deps example_bash_operator runme_0 --logical-date "{date_param}"', 'tasks states-for-dag-run example_bash_operator "manual__{date_param}"', 'tasks states-for-dag-run example_bash_operator --logical-date "{date_param}"', 'tasks clear example_bash_operator --dag-run-id "manual__{date_param}" --task-ids runme_0 -o json', # Task Instances commands + 'taskinstances get example_bash_operator "manual__{date_param}" runme_0', + 'taskinstances get-dependencies example_bash_operator "manual__{date_param}" runme_0', 'taskinstances list example_bash_operator "manual__{date_param}"', # XCom commands - need a Dag run with completed tasks 'xcom add example_bash_operator "manual__{date_param}" runme_0 {xcom_key} \'{{"test": "value"}}\'', diff --git a/airflow-ctl/docs/images/command_hashes.txt b/airflow-ctl/docs/images/command_hashes.txt index 74b5ede7ce3a2..0b34d439f712c 100644 --- a/airflow-ctl/docs/images/command_hashes.txt +++ b/airflow-ctl/docs/images/command_hashes.txt @@ -9,8 +9,8 @@ dagrun:07035226eaaff0a557d3ad9aab4b41b5 jobs:a5b644c5da8889443bb40ee10b599270 pools:19efe105b9515ab1926ebcaf0e028d71 providers:34502fe09dc0b8b0a13e7e46efdffda6 -taskinstances:7e323968c0b585287c2a4ab4339aee1f -tasks:e2ed0ab7f67ebacf4be87b5f396ff782 +taskinstances:bea84117114c2438eb7e7026f6bb7042 +tasks:ea587dc805cadbce81cd640d357f2661 variables:f8fc76d3d398b2780f4e97f7cd816646 version:31f4efdf8de0dbaaa4fac71ff7efecc3 plugins:4864fd8f356704bd2b3cd1aec3567e35 diff --git a/airflow-ctl/docs/images/output_taskinstances.svg b/airflow-ctl/docs/images/output_taskinstances.svg index 9d9e86568e16d..ccde77e2bb834 100644 --- a/airflow-ctl/docs/images/output_taskinstances.svg +++ b/airflow-ctl/docs/images/output_taskinstances.svg @@ -1,4 +1,4 @@ - + - - + + - + - + - + - + - + - + - + - + - + + + + + + + + + + - + - + - - Usage:airflowctl taskinstances [-h] COMMAND... - -Perform TaskInstances operations - -Positional Arguments: -COMMAND -listList all task instances for a given Dag run - -Options: --h, --helpshow this help message and exit + + Usage:airflowctl taskinstances [-h] COMMAND... + +Perform TaskInstances operations + +Positional Arguments: +COMMAND +getGet a task instance for a given Dag run +get-dependencies +Get unmet scheduler dependencies for a task instance +listList all task instances for a given Dag run + +Options: +-h, --helpshow this help message and exit diff --git a/airflow-ctl/docs/images/output_tasks.svg b/airflow-ctl/docs/images/output_tasks.svg index f3fa3c4d94506..9f30d658f385e 100644 --- a/airflow-ctl/docs/images/output_tasks.svg +++ b/airflow-ctl/docs/images/output_tasks.svg @@ -1,4 +1,4 @@ - + - - + + - + - + - + - + - + - + - + - + - + - + - + + + + - + - + - - Usage:airflowctl tasks [-h] COMMAND... - -Perform Tasks operations - -Positional Arguments: -COMMAND -clearClear task instances of a Dag by its ID -states-for-dag-run -Get the status of all task instances in a Dag run - -Options: --h, --helpshow this help message and exit + + Usage:airflowctl tasks [-h] COMMAND... + +Perform Tasks operations + +Positional Arguments: +COMMAND +clearClear task instances of a Dag by its ID +failed-depsReturns the unmet dependencies for a task instance +states-for-dag-run +Get the status of all task instances in a Dag run + +Options: +-h, --helpshow this help message and exit diff --git a/airflow-ctl/src/airflowctl/api/operations.py b/airflow-ctl/src/airflowctl/api/operations.py index 5684dfac8403f..c83b7a724bd6e 100644 --- a/airflow-ctl/src/airflowctl/api/operations.py +++ b/airflow-ctl/src/airflowctl/api/operations.py @@ -69,7 +69,9 @@ ProviderCollectionResponse, QueuedEventCollectionResponse, QueuedEventResponse, + TaskDependencyCollectionResponse, TaskInstanceCollectionResponse, + TaskInstanceResponse, TriggerDAGRunPostBody, VariableBody, VariableCollectionResponse, @@ -777,9 +779,55 @@ def list(self) -> ProviderCollectionResponse | ServerResponseError: return super().execute_list(path="providers", data_model=ProviderCollectionResponse) +def _build_task_instance_path(dag_id: str, dag_run_id: str, task_id: str, map_index: int | None) -> str: + """Build the task instance API path, addressing a mapped task instance when map_index is given.""" + path = f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}" + if map_index is not None and map_index >= 0: + path = f"{path}/{map_index}" + return path + + class TaskInstancesOperations(BaseOperations): """Task instance operations.""" + def get( + self, + dag_id: str, + dag_run_id: str, + task_id: str, + map_index: int | None = None, + *, + suppress_error_log: bool = False, + ) -> TaskInstanceResponse | ServerResponseError: + """Get a task instance for a Dag run.""" + path = _build_task_instance_path( + dag_id=dag_id, dag_run_id=dag_run_id, task_id=task_id, map_index=map_index + ) + self.response = self.client.get( + path, + extensions={"airflowctl_suppress_error_log": suppress_error_log}, + ) + return TaskInstanceResponse.model_validate_json(self.response.content) + + def get_dependencies( + self, + dag_id: str, + dag_run_id: str, + task_id: str, + map_index: int | None = None, + *, + suppress_error_log: bool = False, + ) -> TaskDependencyCollectionResponse | ServerResponseError: + """Get unmet scheduler dependencies for a task instance.""" + path = _build_task_instance_path( + dag_id=dag_id, dag_run_id=dag_run_id, task_id=task_id, map_index=map_index + ) + self.response = self.client.get( + f"{path}/dependencies", + extensions={"airflowctl_suppress_error_log": suppress_error_log}, + ) + return TaskDependencyCollectionResponse.model_validate_json(self.response.content) + def list(self, dag_id: str, dag_run_id: str) -> TaskInstanceCollectionResponse | ServerResponseError: """List task instances for a Dag run.""" return super().execute_list( diff --git a/airflow-ctl/src/airflowctl/ctl/cli_config.py b/airflow-ctl/src/airflowctl/ctl/cli_config.py index 899067d341040..339d480842f55 100755 --- a/airflow-ctl/src/airflowctl/ctl/cli_config.py +++ b/airflow-ctl/src/airflowctl/ctl/cli_config.py @@ -294,6 +294,11 @@ def _load_help_texts_yaml() -> dict[str, dict[str, str]]: ) # Task Commands Args +ARG_TASK_ID = Arg( + flags=("task_id",), + type=str, + help="The task ID", +) ARG_RUN_ID = Arg( flags=("run_id",), type=str, @@ -305,6 +310,12 @@ def _load_help_texts_yaml() -> dict[str, dict[str, str]]: type=str, help="The logical date of the Dag run with a timezone offset (pass this or run_id, not both)", ) +ARG_MAP_INDEX = Arg( + flags=("--map-index",), + type=int, + default=-1, + help="Mapped task index", +) ARG_ACTION_ON_EXISTING_KEY = Arg( flags=("-a", "--action-on-existing-key"), @@ -1085,6 +1096,23 @@ def merge_commands( ) TASK_COMMANDS = ( + ActionCommand( + name="failed-deps", + help="Returns the unmet dependencies for a task instance", + description=( + "Returns the unmet dependencies for a task instance from the perspective of the scheduler. " + "In other words, why a task instance doesn't get scheduled and then queued by the scheduler, " + "and then run by an executor." + ), + func=lazy_load_command("airflowctl.ctl.commands.task_command.failed_deps"), + args=( + ARG_DAG_ID, + ARG_TASK_ID, + ARG_RUN_ID, + ARG_LOGICAL_DATE, + ARG_MAP_INDEX, + ), + ), ActionCommand( name="states-for-dag-run", help="Get the status of all task instances in a Dag run", diff --git a/airflow-ctl/src/airflowctl/ctl/commands/task_command.py b/airflow-ctl/src/airflowctl/ctl/commands/task_command.py index 5881a1bd55675..abd61b23dde99 100644 --- a/airflow-ctl/src/airflowctl/ctl/commands/task_command.py +++ b/airflow-ctl/src/airflowctl/ctl/commands/task_command.py @@ -24,6 +24,7 @@ import rich from airflowctl.api.client import NEW_API_CLIENT, ClientKind, ServerResponseError, provide_api_client +from airflowctl.api.datamodels.generated import TaskInstanceState from airflowctl.ctl.console_formatting import AirflowConsole if TYPE_CHECKING: @@ -74,6 +75,57 @@ def _format_task_instance(ti: TaskInstanceResponse, has_mapped_instances: bool) return data +@provide_api_client(kind=ClientKind.CLI) +def failed_deps(args, api_client=NEW_API_CLIENT) -> None: + """Get task instance dependencies that were not met, from the scheduler's perspective.""" + if (args.run_id is None) == (args.logical_date is None): + rich.print("[red]Provide either run_id or --logical-date, but not both[/red]") + sys.exit(1) + + run_id = args.run_id or _find_run_id_by_logical_date(api_client, args.dag_id, args.logical_date) + + try: + response = api_client.task_instances.get_dependencies( + dag_id=args.dag_id, + dag_run_id=run_id, + task_id=args.task_id, + map_index=args.map_index, + suppress_error_log=True, + ) + if response.dependencies: + print("Task instance dependencies not met:") + for dep in response.dependencies: + print(f"{dep.name}: {dep.reason}") + return + + # The API server evaluates scheduler dependencies only for task instances that have not been + # queued yet, so an empty response for any other state does not mean the dependencies are met. + task_instance = api_client.task_instances.get( + dag_id=args.dag_id, + dag_run_id=run_id, + task_id=args.task_id, + map_index=args.map_index, + suppress_error_log=True, + ) + except ServerResponseError as e: + if e.response.status_code == 404: + map_index_part = f" with map index {args.map_index}" if args.map_index >= 0 else "" + rich.print( + f"[red]Task instance for task {args.task_id!r}{map_index_part} in Dag run " + f"{run_id!r} of Dag {args.dag_id!r} not found[/red]" + ) + sys.exit(1) + raise + + if task_instance.state is None or task_instance.state == TaskInstanceState.SCHEDULED: + print("Task instance dependencies are all met.") + else: + print( + f"Task instance is in the '{task_instance.state.value}' state; scheduler dependencies " + "are only evaluated for task instances that have not yet been queued." + ) + + @provide_api_client(kind=ClientKind.CLI) def states_for_dag_run(args, api_client=NEW_API_CLIENT) -> None: """Get the status of all task instances in a Dag run.""" diff --git a/airflow-ctl/src/airflowctl/ctl/help_texts.yaml b/airflow-ctl/src/airflowctl/ctl/help_texts.yaml index 9c2c591417b14..57ebde04ea9c1 100644 --- a/airflow-ctl/src/airflowctl/ctl/help_texts.yaml +++ b/airflow-ctl/src/airflowctl/ctl/help_texts.yaml @@ -87,6 +87,8 @@ providers: taskinstances: list: "List all task instances for a given Dag run" + get: "Get a task instance for a given Dag run" + get-dependencies: "Get unmet scheduler dependencies for a task instance" tasks: clear: "Clear task instances of a Dag by its ID" diff --git a/airflow-ctl/tests/airflow_ctl/api/test_operations.py b/airflow-ctl/tests/airflow_ctl/api/test_operations.py index ec6695077c8b2..eb3724f848b33 100644 --- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py +++ b/airflow-ctl/tests/airflow_ctl/api/test_operations.py @@ -92,6 +92,8 @@ QueuedEventCollectionResponse, QueuedEventResponse, ReprocessBehavior, + TaskDependencyCollectionResponse, + TaskDependencyResponse, TaskInstanceCollectionResponse, TaskInstanceResponse, TaskInstanceState, @@ -1596,15 +1598,61 @@ class TestTaskInstancesOperations: task_instances=[task_instance_response], total_entries=1, ) + task_dependency_collection_response = TaskDependencyCollectionResponse( + dependencies=[TaskDependencyResponse(name="Trigger Rule", reason="upstream tasks not done")], + ) - def test_list(self): + def _make_client_asserting_path(self, expected_path: str, response_model) -> Client: def handle_request(request: httpx.Request) -> httpx.Response: - assert request.url.path == "/api/v2/dags/dag_id/dagRuns/dag_run_id/taskInstances" - return httpx.Response( - 200, json=json.loads(self.task_instance_collection_response.model_dump_json()) - ) + assert request.url.path == expected_path + return httpx.Response(200, json=json.loads(response_model.model_dump_json())) - client = make_api_client(transport=httpx.MockTransport(handle_request)) + return make_api_client(transport=httpx.MockTransport(handle_request)) + + @pytest.mark.parametrize("map_index", [None, -1]) + def test_get(self, map_index): + client = self._make_client_asserting_path( + "/api/v2/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id", self.task_instance_response + ) + response = client.task_instances.get( + dag_id="dag_id", dag_run_id="dag_run_id", task_id="task_id", map_index=map_index + ) + assert response == self.task_instance_response + + def test_get_with_map_index(self): + client = self._make_client_asserting_path( + "/api/v2/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/3", self.task_instance_response + ) + response = client.task_instances.get( + dag_id="dag_id", dag_run_id="dag_run_id", task_id="task_id", map_index=3 + ) + assert response == self.task_instance_response + + @pytest.mark.parametrize("map_index", [None, -1]) + def test_get_dependencies(self, map_index): + client = self._make_client_asserting_path( + "/api/v2/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/dependencies", + self.task_dependency_collection_response, + ) + response = client.task_instances.get_dependencies( + dag_id="dag_id", dag_run_id="dag_run_id", task_id="task_id", map_index=map_index + ) + assert response == self.task_dependency_collection_response + + def test_get_dependencies_with_map_index(self): + client = self._make_client_asserting_path( + "/api/v2/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/3/dependencies", + self.task_dependency_collection_response, + ) + response = client.task_instances.get_dependencies( + dag_id="dag_id", dag_run_id="dag_run_id", task_id="task_id", map_index=3 + ) + assert response == self.task_dependency_collection_response + + def test_list(self): + client = self._make_client_asserting_path( + "/api/v2/dags/dag_id/dagRuns/dag_run_id/taskInstances", self.task_instance_collection_response + ) response = client.task_instances.list(dag_id="dag_id", dag_run_id="dag_run_id") assert response == self.task_instance_collection_response diff --git a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py index cb134726711b2..cd445a6e41a7b 100644 --- a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py +++ b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py @@ -24,6 +24,8 @@ import pytest from airflowctl.api.datamodels.generated import ( + TaskDependencyCollectionResponse, + TaskDependencyResponse, TaskInstanceCollectionResponse, TaskInstanceResponse, TaskInstanceState, @@ -43,6 +45,318 @@ def _normalize_rich_output(text: str) -> str: return " ".join(text.split()) +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.list.return_value.dag_runs = [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() + assert capsys.readouterr().out == ( + "Task instance dependencies not met:\n" + "Dagrun Running: Task instance's dagrun was not in the 'running' state\n" + "Trigger Rule: requires all upstream tasks to have succeeded\n" + ) + + @pytest.mark.parametrize("state", [None, TaskInstanceState.SCHEDULED]) + def test_failed_deps_all_met(self, state, capsys): + api_client = self._make_api_client(state=state) + + 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.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, + ) + assert capsys.readouterr().out == "Task instance dependencies are all met.\n" + + def test_failed_deps_not_evaluated_for_already_queued_task(self, capsys): + api_client = self._make_api_client(state=TaskInstanceState.SUCCESS) + + task_command.failed_deps( + self.parser.parse_args(["tasks", "failed-deps", self.dag_id, self.task_id, self.run_id]), + api_client=api_client, + ) + + assert capsys.readouterr().out == ( + "Task instance is in the 'success' state; scheduler dependencies are only evaluated " + "for task instances that have not yet been queued.\n" + ) + + def test_failed_deps_with_map_index(self, capsys): + api_client = self._make_api_client( + dependencies=[TaskDependencyResponse(name="Pool Slots Available", reason="pool is full")] + ) + + task_command.failed_deps( + self.parser.parse_args( + ["tasks", "failed-deps", self.dag_id, self.task_id, self.run_id, "--map-index", "3"] + ), + 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=3, + suppress_error_log=True, + ) + assert capsys.readouterr().out == ( + "Task instance dependencies not met:\nPool Slots Available: pool is full\n" + ) + + def test_failed_deps_by_logical_date(self, capsys): + api_client = self._make_api_client( + dependencies=[TaskDependencyResponse(name="Trigger Rule", reason="upstream tasks not done")] + ) + + task_command.failed_deps( + self.parser.parse_args( + [ + "tasks", + "failed-deps", + self.dag_id, + self.task_id, + "--logical-date", + self.logical_date.isoformat(), + ] + ), + api_client=api_client, + ) + + api_client.dag_runs.list.assert_called_once_with( + dag_id=self.dag_id, + logical_date_gte=self.logical_date, + logical_date_lte=self.logical_date, + order_by="-id", + limit=1, + suppress_error_log=True, + ) + 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, + ) + assert capsys.readouterr().out == ( + "Task instance dependencies not met:\nTrigger Rule: upstream tasks not done\n" + ) + + @pytest.mark.parametrize( + "extra_args", + [ + [], + ["test_run", "--logical-date", "2025-01-01T00:00:00+00:00"], + ], + ids=["neither", "both"], + ) + def test_failed_deps_requires_exactly_one_of_run_id_and_logical_date(self, extra_args, capsys): + api_client = self._make_api_client() + + with pytest.raises(SystemExit, match="1"): + task_command.failed_deps( + self.parser.parse_args(["tasks", "failed-deps", self.dag_id, self.task_id, *extra_args]), + api_client=api_client, + ) + + api_client.task_instances.get_dependencies.assert_not_called() + assert _normalize_rich_output(capsys.readouterr().out) == ( + "Provide either run_id or --logical-date, but not both" + ) + + @pytest.mark.parametrize( + ("logical_date", "expected_message"), + [ + ("not-a-date", "Invalid --logical-date: 'not-a-date'"), + ("2025-01-01T00:00:00", "--logical-date must include a timezone offset"), + ], + ids=["unparsable", "naive"], + ) + def test_failed_deps_rejects_bad_logical_date(self, logical_date, expected_message, capsys): + api_client = self._make_api_client() + + with pytest.raises(SystemExit, match="1"): + task_command.failed_deps( + self.parser.parse_args( + ["tasks", "failed-deps", self.dag_id, self.task_id, "--logical-date", logical_date] + ), + api_client=api_client, + ) + + api_client.dag_runs.list.assert_not_called() + assert _normalize_rich_output(capsys.readouterr().out) == expected_message + + @pytest.mark.parametrize("list_failure", ["no_matching_run", "dag_not_found_404"]) + def test_failed_deps_dag_run_not_found_by_logical_date(self, list_failure, capsys): + api_client = self._make_api_client() + if list_failure == "no_matching_run": + api_client.dag_runs.list.return_value.dag_runs = [] + else: + api_client.dag_runs.list.side_effect = _make_server_error(404) + + with pytest.raises(SystemExit, match="1"): + task_command.failed_deps( + self.parser.parse_args( + [ + "tasks", + "failed-deps", + self.dag_id, + self.task_id, + "--logical-date", + self.logical_date.isoformat(), + ] + ), + api_client=api_client, + ) + + api_client.task_instances.get_dependencies.assert_not_called() + assert _normalize_rich_output(capsys.readouterr().out) == ( + "Dag run for test_dag with logical date '2025-01-01T00:00:00+00:00' not found" + ) + + @pytest.mark.parametrize( + ("extra_args", "expected_message"), + [ + ([], "Task instance for task 'test_task' in Dag run 'test_run' of Dag 'test_dag' not found"), + ( + ["--map-index", "3"], + "Task instance for task 'test_task' with map index 3 in Dag run 'test_run' " + "of Dag 'test_dag' not found", + ), + ], + ) + def test_failed_deps_task_instance_not_found(self, extra_args, expected_message, capsys): + api_client = self._make_api_client() + api_client.task_instances.get_dependencies.side_effect = _make_server_error(404) + + with pytest.raises(SystemExit, match="1"): + task_command.failed_deps( + self.parser.parse_args( + ["tasks", "failed-deps", self.dag_id, self.task_id, self.run_id, *extra_args] + ), + api_client=api_client, + ) + + assert _normalize_rich_output(capsys.readouterr().out) == expected_message + + def test_failed_deps_task_instance_deleted_between_requests(self, capsys): + api_client = self._make_api_client() + api_client.task_instances.get.side_effect = _make_server_error(404) + + with pytest.raises(SystemExit, match="1"): + task_command.failed_deps( + self.parser.parse_args(["tasks", "failed-deps", self.dag_id, self.task_id, self.run_id]), + api_client=api_client, + ) + + assert _normalize_rich_output(capsys.readouterr().out) == ( + "Task instance for task 'test_task' in Dag run 'test_run' of Dag 'test_dag' not found" + ) + + def test_failed_deps_reraises_non_404_task_instance_get_error(self): + api_client = self._make_api_client() + api_client.task_instances.get.side_effect = _make_server_error(500) + + with pytest.raises(ServerResponseError): + task_command.failed_deps( + self.parser.parse_args(["tasks", "failed-deps", self.dag_id, self.task_id, self.run_id]), + api_client=api_client, + ) + + def test_failed_deps_reraises_non_404_error(self): + api_client = self._make_api_client() + api_client.task_instances.get_dependencies.side_effect = _make_server_error(500) + + with pytest.raises(ServerResponseError): + task_command.failed_deps( + self.parser.parse_args(["tasks", "failed-deps", self.dag_id, self.task_id, self.run_id]), + api_client=api_client, + ) + + def test_failed_deps_reraises_non_404_dag_run_list_error(self): + api_client = self._make_api_client() + api_client.dag_runs.list.side_effect = _make_server_error(500) + + with pytest.raises(ServerResponseError): + task_command.failed_deps( + self.parser.parse_args( + [ + "tasks", + "failed-deps", + self.dag_id, + self.task_id, + "--logical-date", + self.logical_date.isoformat(), + ] + ), + api_client=api_client, + ) + + class TestStatesForDagRun: parser = cli_parser.get_parser() dag_id = "test_dag" diff --git a/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py b/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py index d568cacc0981c..fcbd3d748b2f2 100644 --- a/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py +++ b/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py @@ -833,6 +833,8 @@ def test_apply_datamodel_defaults_clear_task_instances_invalid_json_task_ids(sel ("assets", "get", "Retrieve an asset by its ID"), ("connections", "get", "Retrieve a connection by its ID"), ("taskinstances", "list", "List all task instances for a given Dag run"), + ("taskinstances", "get", "Get a task instance for a given Dag run"), + ("taskinstances", "get-dependencies", "Get unmet scheduler dependencies for a task instance"), ("tasks", "clear", "Clear task instances of a Dag by its ID"), ], )