diff --git a/airflow-core/src/airflow/cli/commands/task_command.py b/airflow-core/src/airflow/cli/commands/task_command.py index af944526184b0..a39f07e442dfe 100644 --- a/airflow-core/src/airflow/cli/commands/task_command.py +++ b/airflow-core/src/airflow/cli/commands/task_command.py @@ -30,7 +30,7 @@ from airflow import settings from airflow._shared.timezones import timezone from airflow.cli.simple_table import AirflowConsole -from airflow.cli.utils import fetch_dag_run_from_run_id_or_logical_date_string +from airflow.cli.utils import deprecated_for_airflowctl, fetch_dag_run_from_run_id_or_logical_date_string from airflow.exceptions import AirflowConfigException, DagRunNotFound, NotMapped, TaskInstanceNotFound from airflow.models import TaskInstance from airflow.models.dag_version import DagVersion @@ -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 20fc203595301..d23f1d8afb68d 100644 --- a/airflow-core/tests/unit/cli/commands/test_command_deprecations.py +++ b/airflow-core/tests/unit/cli/commands/test_command_deprecations.py @@ -37,6 +37,7 @@ dag_command, pool_command, provider_command, + task_command, variable_command, ) @@ -74,6 +75,7 @@ (provider_command.providers_list, "airflowctl providers list"), (config_command.get_value, "airflowctl config get"), (config_command.show_config, "airflowctl config list"), + (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 924cbde90d598..93b9e09cd493a 100644 --- a/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py +++ b/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py @@ -96,6 +96,12 @@ def date_param(): "dags update example_bash_operator --no-is-paused", # 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}"', + # Task Instances commands + 'taskinstances get example_bash_operator "manual__{date_param}" runme_0', + 'taskinstances get-dependencies example_bash_operator "manual__{date_param}" runme_0', # XCom commands - need a Dag run with completed tasks 'xcom add example_bash_operator "manual__{date_param}" runme_0 {xcom_key} \'{{"test": "value"}}\'', 'xcom get example_bash_operator "manual__{date_param}" runme_0 {xcom_key}', diff --git a/airflow-ctl/docs/conf.py b/airflow-ctl/docs/conf.py index 055de186d3ef2..6580a8cbf0c51 100644 --- a/airflow-ctl/docs/conf.py +++ b/airflow-ctl/docs/conf.py @@ -283,7 +283,10 @@ def add_airflow_ctl_exclude_patterns_to_sphinx(exclude_patterns: list[str]): redirects_file = "redirects.txt" # -- Options for sphinxcontrib-spelling ---------------------------------------- -spelling_word_list_filename = [SPELLING_WORDLIST_PATH.as_posix()] +spelling_word_list_filename = [ + SPELLING_WORDLIST_PATH.as_posix(), + (pathlib.Path(__file__).parent / "spelling_wordlist.txt").as_posix(), +] spelling_exclude_patterns = ["project.rst", "changelog.rst"] spelling_ignore_contributor_names = False diff --git a/airflow-ctl/docs/howto/index.rst b/airflow-ctl/docs/howto/index.rst index 59f806d4f44b3..92f9244321550 100644 --- a/airflow-ctl/docs/howto/index.rst +++ b/airflow-ctl/docs/howto/index.rst @@ -213,6 +213,20 @@ These visual references show the full command syntax, options, and parameters fo :width: 60% :alt: airflowctl Providers Command +**Task Instances** +'''''''''''''''''' +.. image:: ../images/output_taskinstances.svg + :target: https://raw.githubusercontent.com/apache/airflow/main/airflow-ctl/docs/images/output_taskinstances.svg + :width: 60% + :alt: airflowctl Task Instances Command + +**Tasks** +''''''''' +.. image:: ../images/output_tasks.svg + :target: https://raw.githubusercontent.com/apache/airflow/main/airflow-ctl/docs/images/output_tasks.svg + :width: 60% + :alt: airflowctl Tasks Command + **Variables** ''''''''''''' .. image:: ../images/output_variables.svg diff --git a/airflow-ctl/docs/images/command_hashes.txt b/airflow-ctl/docs/images/command_hashes.txt index 36eec17c2eb9f..c38963edb7f05 100644 --- a/airflow-ctl/docs/images/command_hashes.txt +++ b/airflow-ctl/docs/images/command_hashes.txt @@ -1,4 +1,4 @@ -main:27a22c00dcf32e7a1a4f06672dc8e3c8 +main:0460d9c03248bee26207b20b05aa36b9 assets:6419e20452692f577c4c6f570b74be0c auth:d79e9c7d00c432bdbcbc2a86e2e32053 backfill:74c8737b0a62a86ed3605fa9e6165874 @@ -9,6 +9,8 @@ dagrun:07035226eaaff0a557d3ad9aab4b41b5 jobs:a5b644c5da8889443bb40ee10b599270 pools:19efe105b9515ab1926ebcaf0e028d71 providers:34502fe09dc0b8b0a13e7e46efdffda6 +taskinstances:828543543b03adbcd743a6f4eda308d8 +tasks:33ee6ea468f816f71f1c363b785dbb77 variables:f8fc76d3d398b2780f4e97f7cd816646 version:31f4efdf8de0dbaaa4fac71ff7efecc3 plugins:4864fd8f356704bd2b3cd1aec3567e35 diff --git a/airflow-ctl/docs/images/output_main.svg b/airflow-ctl/docs/images/output_main.svg index f586877bce8eb..9c8cec5269b68 100644 --- a/airflow-ctl/docs/images/output_main.svg +++ b/airflow-ctl/docs/images/output_main.svg @@ -1,4 +1,4 @@ - + - - + + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + + + + + + + + + + - + - + - - Usage:airflowctl [-h] GROUP_OR_COMMAND... - -Positional Arguments: -GROUP_OR_COMMAND - -    Groups -assetsPerform Assets operations -authManage authentication for CLI. Either pass token from -environment variable/parameter or pass username and -password. -backfillPerform Backfill operations -configPerform Config operations -connectionsPerform Connections operations -dagrunPerform DagRun operations -dagsPerform Dags operations -jobsPerform Jobs operations -pluginsPerform Plugins operations -poolsPerform Pools operations -providersPerform Providers operations -variablesPerform Variables operations -xcomPerform XCom operations - -    Commands: -versionShow version information - -Options: --h, --helpshow this help message and exit + + Usage:airflowctl [-h] GROUP_OR_COMMAND... + +Positional Arguments: +GROUP_OR_COMMAND + +    Groups +assetsPerform Assets operations +authManage authentication for CLI. Either pass token from +environment variable/parameter or pass username and +password. +backfillPerform Backfill operations +configPerform Config operations +connectionsPerform Connections operations +dagrunPerform DagRun operations +dagsPerform Dags operations +jobsPerform Jobs operations +pluginsPerform Plugins operations +poolsPerform Pools operations +providersPerform Providers operations +taskinstances +Perform TaskInstances operations +tasksManage Airflow tasks +variablesPerform Variables operations +xcomPerform XCom operations + +    Commands: +versionShow version information + +Options: +-h, --helpshow this help message and exit diff --git a/airflow-ctl/docs/images/output_taskinstances.svg b/airflow-ctl/docs/images/output_taskinstances.svg new file mode 100644 index 0000000000000..4ecbab5b13592 --- /dev/null +++ b/airflow-ctl/docs/images/output_taskinstances.svg @@ -0,0 +1,105 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + 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 + +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 new file mode 100644 index 0000000000000..6df6ad9af71c6 --- /dev/null +++ b/airflow-ctl/docs/images/output_tasks.svg @@ -0,0 +1,101 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + Usage:airflowctl tasks [-h] COMMAND... + +Manage Airflow tasks + +Positional Arguments: +COMMAND +failed-deps +Returns the unmet dependencies for a task instance + +Options: +-h, --helpshow this help message and exit + + + + diff --git a/airflow-ctl/docs/spelling_wordlist.txt b/airflow-ctl/docs/spelling_wordlist.txt new file mode 100644 index 0000000000000..04c8494c855eb --- /dev/null +++ b/airflow-ctl/docs/spelling_wordlist.txt @@ -0,0 +1 @@ +taskinstances diff --git a/airflow-ctl/src/airflowctl/api/client.py b/airflow-ctl/src/airflowctl/api/client.py index de2cec4ea5f33..7baf87684dd6c 100644 --- a/airflow-ctl/src/airflowctl/api/client.py +++ b/airflow-ctl/src/airflowctl/api/client.py @@ -57,6 +57,7 @@ PoolsOperations, ProvidersOperations, ServerResponseError, + TaskInstancesOperations, VariablesOperations, VersionOperations, XComOperations, @@ -450,6 +451,12 @@ def providers(self): """Operations related to providers.""" return ProvidersOperations(self) + @lru_cache() # type: ignore[prop-decorator] + @property + def task_instances(self): + """Operations related to task instances.""" + return TaskInstancesOperations(self) + @lru_cache() # type: ignore[prop-decorator] @property def variables(self): diff --git a/airflow-ctl/src/airflowctl/api/operations.py b/airflow-ctl/src/airflowctl/api/operations.py index f58ad78adeb37..836e1506d3637 100644 --- a/airflow-ctl/src/airflowctl/api/operations.py +++ b/airflow-ctl/src/airflowctl/api/operations.py @@ -68,6 +68,8 @@ ProviderCollectionResponse, QueuedEventCollectionResponse, QueuedEventResponse, + TaskDependencyCollectionResponse, + TaskInstanceResponse, TriggerDAGRunPostBody, VariableBody, VariableCollectionResponse, @@ -634,6 +636,8 @@ def list( logical_date_gte: datetime.datetime | None = None, logical_date_lte: datetime.datetime | None = None, order_by: str | None = None, + *, + suppress_error_log: bool = False, ) -> DAGRunCollectionResponse | ServerResponseError: """ List dag runs (at most `limit` results). @@ -647,6 +651,7 @@ def list( logical_date_gte: Filter dag runs with a logical date greater than or equal to this value. logical_date_lte: Filter dag runs with a logical date less than or equal to this value. order_by: Order the results by the specified field. + suppress_error_log: Skip client-side error logging, for callers handling the error themselves. """ # Use "~" for all DAGs if dag_id is not specified if not dag_id: @@ -663,7 +668,11 @@ def list( ) try: - self.response = self.client.get(f"/dags/{dag_id}/dagRuns", params=params) + self.response = self.client.get( + f"/dags/{dag_id}/dagRuns", + params=params, + extensions={"airflowctl_suppress_error_log": suppress_error_log}, + ) return DAGRunCollectionResponse.model_validate_json(self.response.content) except ServerResponseError as e: raise e @@ -750,6 +759,62 @@ 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 + ) + try: + self.response = self.client.get( + path, + extensions={"airflowctl_suppress_error_log": suppress_error_log}, + ) + return TaskInstanceResponse.model_validate_json(self.response.content) + except ServerResponseError as e: + raise e + + 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 + ) + try: + self.response = self.client.get( + f"{path}/dependencies", + extensions={"airflowctl_suppress_error_log": suppress_error_log}, + ) + return TaskDependencyCollectionResponse.model_validate_json(self.response.content) + except ServerResponseError as e: + raise e + + class VariablesOperations(BaseOperations): """Variable operations.""" diff --git a/airflow-ctl/src/airflowctl/ctl/cli_config.py b/airflow-ctl/src/airflowctl/ctl/cli_config.py index 40ee7a1c67a21..1c1af9c255a46 100755 --- a/airflow-ctl/src/airflowctl/ctl/cli_config.py +++ b/airflow-ctl/src/airflowctl/ctl/cli_config.py @@ -279,7 +279,7 @@ def _load_help_texts_yaml() -> dict[str, dict[str, str]]: ARG_DAG_ID = Arg( flags=("dag_id",), type=str, - help="The Dag ID of the Dag to pause or unpause", + help="The Dag ID", ) ARG_LOGICAL_DATE_OR_RUN_ID = Arg( flags=("logical_date_or_run_id",), @@ -287,6 +287,30 @@ def _load_help_texts_yaml() -> dict[str, dict[str, str]]: help="The logical date with a timezone offset or run ID of the Dag run", ) +# 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, + nargs="?", + help="The run ID of the Dag run (pass this or --logical-date, not both)", +) +ARG_LOGICAL_DATE = Arg( + flags=("--logical-date",), + 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"), type=str, @@ -1034,6 +1058,26 @@ 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, + ), + ), +) + VARIABLE_COMMANDS = ( ActionCommand( name="import", @@ -1070,6 +1114,11 @@ def merge_commands( help="Manage Airflow pools", subcommands=POOL_COMMANDS, ), + GroupCommand( + name="tasks", + help="Manage Airflow tasks", + subcommands=TASK_COMMANDS, + ), ActionCommand( name="version", help="Show version information", diff --git a/airflow-ctl/src/airflowctl/ctl/commands/task_command.py b/airflow-ctl/src/airflowctl/ctl/commands/task_command.py new file mode 100644 index 0000000000000..bc3ebda4a81f4 --- /dev/null +++ b/airflow-ctl/src/airflowctl/ctl/commands/task_command.py @@ -0,0 +1,107 @@ +# 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 sys + +import rich + +from airflowctl.api.client import NEW_API_CLIENT, ClientKind, ServerResponseError, provide_api_client +from airflowctl.api.datamodels.generated import TaskInstanceState + + +def _find_run_id_by_logical_date(api_client, dag_id: str, value: str) -> str: + """Find the run ID of the Dag run with an exact logical date match.""" + try: + logical_date = datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + rich.print(f"[red]Invalid --logical-date: {value!r}[/red]") + sys.exit(1) + if logical_date.tzinfo is None: + rich.print("[red]--logical-date must include a timezone offset[/red]") + sys.exit(1) + + dag_runs = [] + try: + dag_runs = api_client.dag_runs.list( + dag_id=dag_id, + logical_date_gte=logical_date, + logical_date_lte=logical_date, + order_by="-id", + limit=1, + suppress_error_log=True, + ).dag_runs + except ServerResponseError as e: + if e.response.status_code != 404: + raise + if not dag_runs: + rich.print(f"[red]Dag run for {dag_id} with logical date {value!r} not found[/red]") + sys.exit(1) + return dag_runs[0].dag_run_id + + +@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." + ) diff --git a/airflow-ctl/src/airflowctl/ctl/help_texts.yaml b/airflow-ctl/src/airflowctl/ctl/help_texts.yaml index d0ebdcfeaebdd..68f5377661f84 100644 --- a/airflow-ctl/src/airflowctl/ctl/help_texts.yaml +++ b/airflow-ctl/src/airflowctl/ctl/help_texts.yaml @@ -85,6 +85,10 @@ pools: providers: list: "List all installed Airflow providers" +taskinstances: + get: "Get a task instance for a given Dag run" + get-dependencies: "Get unmet scheduler dependencies for a task instance" + variables: get: "Retrieve a variable by its key" list: "List all variables" diff --git a/airflow-ctl/tests/airflow_ctl/api/test_operations.py b/airflow-ctl/tests/airflow_ctl/api/test_operations.py index dc5cdd396db8b..0522cbde864bd 100644 --- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py +++ b/airflow-ctl/tests/airflow_ctl/api/test_operations.py @@ -91,6 +91,10 @@ QueuedEventCollectionResponse, QueuedEventResponse, ReprocessBehavior, + TaskDependencyCollectionResponse, + TaskDependencyResponse, + TaskInstanceResponse, + TaskInstanceState, TriggerDAGRunPostBody, VariableBody, VariableCollectionResponse, @@ -1503,6 +1507,75 @@ 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")], + ) + + def _make_client_asserting_path(self, expected_path: str, response_model) -> Client: + def handle_request(request: httpx.Request) -> httpx.Response: + assert request.url.path == expected_path + return httpx.Response(200, json=json.loads(response_model.model_dump_json())) + + 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 + + class TestVariablesOperations: key = "key" value = "val" 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 new file mode 100644 index 0000000000000..c166d1e357062 --- /dev/null +++ b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py @@ -0,0 +1,356 @@ +# 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 _make_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) + + +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, + ) 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 281bd2b040b04..d5458c82ef4ce 100644 --- a/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py +++ b/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py @@ -784,6 +784,8 @@ def test_apply_datamodel_defaults_other_datamodel(self): [ ("assets", "get", "Retrieve an asset by its ID"), ("connections", "get", "Retrieve a connection by its ID"), + ("taskinstances", "get", "Get a task instance for a given Dag run"), + ("taskinstances", "get-dependencies", "Get unmet scheduler dependencies for a task instance"), ], ) def test_help_texts_used_for_auto_generated_commands(self, group_name, subcommand_name, expected_help): @@ -797,3 +799,4 @@ def test_help_texts_used_for_auto_generated_commands(self, group_name, subcomman "Help message should match the help_text.yaml" ) return + pytest.fail(f"Auto-generated command not found: {group_name} {subcommand_name}") diff --git a/scripts/in_container/run_capture_airflowctl_help.py b/scripts/in_container/run_capture_airflowctl_help.py index 9529dbe04390c..52753d6a551eb 100644 --- a/scripts/in_container/run_capture_airflowctl_help.py +++ b/scripts/in_container/run_capture_airflowctl_help.py @@ -45,6 +45,8 @@ "jobs", "pools", "providers", + "taskinstances", + "tasks", "variables", "version", "plugins",