diff --git a/airflow-core/src/airflow/cli/commands/task_command.py b/airflow-core/src/airflow/cli/commands/task_command.py index af944526184b0..e2a1ab28e0495 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 @@ -354,6 +354,7 @@ def _guess_debugger() -> _SupportedDebugger: raise exc +@deprecated_for_airflowctl("airflowctl tasks states-for-dag-run") @cli_utils.action_cli(check_db=False) @suppress_logs_and_warning @providers_configuration_loaded 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..65982d65adb6f 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_states_for_dag_run, "airflowctl tasks states-for-dag-run"), ] 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..87d9851cbf0c0 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,11 @@ 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 states-for-dag-run example_bash_operator "manual__{date_param}"', + 'tasks states-for-dag-run example_bash_operator --logical-date "{date_param}"', + # Task Instances commands + '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"}}\'', '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..10e59c5959f8b 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:7e323968c0b585287c2a4ab4339aee1f +tasks:089b19625c893d189b5fc02b3abd547e 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..9d9e86568e16d --- /dev/null +++ b/airflow-ctl/docs/images/output_taskinstances.svg @@ -0,0 +1,97 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + 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 + + + + diff --git a/airflow-ctl/docs/images/output_tasks.svg b/airflow-ctl/docs/images/output_tasks.svg new file mode 100644 index 0000000000000..d9349c117529b --- /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 +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/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..4707b6e9f4993 100644 --- a/airflow-ctl/src/airflowctl/api/operations.py +++ b/airflow-ctl/src/airflowctl/api/operations.py @@ -68,6 +68,7 @@ ProviderCollectionResponse, QueuedEventCollectionResponse, QueuedEventResponse, + TaskInstanceCollectionResponse, TriggerDAGRunPostBody, VariableBody, VariableCollectionResponse, @@ -634,6 +635,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 +650,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 +667,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 +758,17 @@ def list(self) -> ProviderCollectionResponse | ServerResponseError: return super().execute_list(path="providers", data_model=ProviderCollectionResponse) +class TaskInstancesOperations(BaseOperations): + """Task instance operations.""" + + def list(self, dag_id: str, dag_run_id: str) -> TaskInstanceCollectionResponse | ServerResponseError: + """List task instances for a Dag run.""" + return super().execute_list( + path=f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances", + data_model=TaskInstanceCollectionResponse, + ) + + 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..47270d9c29c10 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,19 @@ 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_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_ACTION_ON_EXISTING_KEY = Arg( flags=("-a", "--action-on-existing-key"), type=str, @@ -1034,6 +1047,24 @@ def merge_commands( ), ) +TASK_COMMANDS = ( + ActionCommand( + name="states-for-dag-run", + help="Get the status of all task instances in a Dag run", + description=( + "Get the status of all task instances in a Dag run. " + "Select the run with either run_id or --logical-date (pass exactly one)." + ), + func=lazy_load_command("airflowctl.ctl.commands.task_command.states_for_dag_run"), + args=( + ARG_DAG_ID, + ARG_RUN_ID, + ARG_LOGICAL_DATE, + ARG_OUTPUT, + ), + ), +) + VARIABLE_COMMANDS = ( ActionCommand( name="import", @@ -1070,6 +1101,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..5881a1bd55675 --- /dev/null +++ b/airflow-ctl/src/airflowctl/ctl/commands/task_command.py @@ -0,0 +1,99 @@ +# 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 +from typing import TYPE_CHECKING + +import rich + +from airflowctl.api.client import NEW_API_CLIENT, ClientKind, ServerResponseError, provide_api_client +from airflowctl.ctl.console_formatting import AirflowConsole + +if TYPE_CHECKING: + from airflowctl.api.datamodels.generated import TaskInstanceResponse + + +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 + + +def _format_task_instance(ti: TaskInstanceResponse, has_mapped_instances: bool) -> dict[str, str]: + data = { + "dag_id": ti.dag_id, + "logical_date": ti.logical_date.isoformat() if ti.logical_date else "", + "task_id": ti.task_id, + "state": ti.state.value if ti.state else "", + "start_date": ti.start_date.isoformat() if ti.start_date else "", + "end_date": ti.end_date.isoformat() if ti.end_date else "", + } + if has_mapped_instances: + data["map_index"] = str(ti.map_index) if ti.map_index >= 0 else "" + return data + + +@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.""" + 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: + task_instances = api_client.task_instances.list(dag_id=args.dag_id, dag_run_id=run_id).task_instances + except ServerResponseError as e: + if e.response.status_code == 404: + rich.print(f"[red]Dag run {run_id!r} of Dag {args.dag_id!r} not found[/red]") + sys.exit(1) + raise + + has_mapped_instances = any(ti.map_index >= 0 for ti in task_instances) + + AirflowConsole().print_as( + data=[_format_task_instance(ti, has_mapped_instances) for ti in task_instances], + output=args.output, + ) diff --git a/airflow-ctl/src/airflowctl/ctl/help_texts.yaml b/airflow-ctl/src/airflowctl/ctl/help_texts.yaml index d0ebdcfeaebdd..1c9b23d3f9491 100644 --- a/airflow-ctl/src/airflowctl/ctl/help_texts.yaml +++ b/airflow-ctl/src/airflowctl/ctl/help_texts.yaml @@ -85,6 +85,9 @@ pools: providers: list: "List all installed Airflow providers" +taskinstances: + list: "List all task instances for a given Dag run" + 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..b4ed5260502f0 100644 --- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py +++ b/airflow-ctl/tests/airflow_ctl/api/test_operations.py @@ -91,6 +91,9 @@ QueuedEventCollectionResponse, QueuedEventResponse, ReprocessBehavior, + TaskInstanceCollectionResponse, + TaskInstanceResponse, + TaskInstanceState, TriggerDAGRunPostBody, VariableBody, VariableCollectionResponse, @@ -1503,6 +1506,40 @@ 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_instance_collection_response = TaskInstanceCollectionResponse( + task_instances=[task_instance_response], + total_entries=1, + ) + + def test_list(self): + 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()) + ) + + client = make_api_client(transport=httpx.MockTransport(handle_request)) + response = client.task_instances.list(dag_id="dag_id", dag_run_id="dag_run_id") + assert response == self.task_instance_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..cb134726711b2 --- /dev/null +++ b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py @@ -0,0 +1,284 @@ +# 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 ( + TaskInstanceCollectionResponse, + 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 TestStatesForDagRun: + parser = cli_parser.get_parser() + dag_id = "test_dag" + run_id = "test_run" + logical_date = datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc) + + def _make_task_instance( + self, + task_id: str, + *, + map_index: int = -1, + state: TaskInstanceState | None = TaskInstanceState.SUCCESS, + start_date: datetime.datetime | None = None, + end_date: datetime.datetime | None = None, + ) -> TaskInstanceResponse: + return TaskInstanceResponse( + id=uuid.uuid4(), + task_id=task_id, + dag_id=self.dag_id, + dag_run_id=self.run_id, + map_index=map_index, + logical_date=self.logical_date, + run_after=self.logical_date, + state=state, + start_date=start_date, + end_date=end_date, + try_number=1, + max_tries=0, + task_display_name=task_id, + dag_display_name=self.dag_id, + pool="default_pool", + pool_slots=1, + executor_config="{}", + ) + + def _make_api_client(self, task_instances: list[TaskInstanceResponse]) -> 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.list.return_value = TaskInstanceCollectionResponse( + task_instances=task_instances, + total_entries=len(task_instances), + ) + return api_client + + @mock.patch("airflowctl.ctl.commands.task_command.AirflowConsole") + def test_states_for_dag_run_by_run_id(self, mock_console_cls): + start_date = datetime.datetime(2025, 1, 1, 1, tzinfo=datetime.timezone.utc) + end_date = datetime.datetime(2025, 1, 1, 2, tzinfo=datetime.timezone.utc) + api_client = self._make_api_client( + [ + self._make_task_instance("task_a", start_date=start_date, end_date=end_date), + self._make_task_instance("task_b", state=None), + ] + ) + + task_command.states_for_dag_run( + self.parser.parse_args(["tasks", "states-for-dag-run", self.dag_id, self.run_id]), + api_client=api_client, + ) + + api_client.dag_runs.list.assert_not_called() + api_client.task_instances.list.assert_called_once_with(dag_id=self.dag_id, dag_run_id=self.run_id) + mock_console_cls.return_value.print_as.assert_called_once_with( + data=[ + { + "dag_id": self.dag_id, + "logical_date": "2025-01-01T00:00:00+00:00", + "task_id": "task_a", + "state": "success", + "start_date": "2025-01-01T01:00:00+00:00", + "end_date": "2025-01-01T02:00:00+00:00", + }, + { + "dag_id": self.dag_id, + "logical_date": "2025-01-01T00:00:00+00:00", + "task_id": "task_b", + "state": "", + "start_date": "", + "end_date": "", + }, + ], + output="json", + ) + + @mock.patch("airflowctl.ctl.commands.task_command.AirflowConsole") + def test_states_for_dag_run_by_logical_date(self, mock_console_cls): + api_client = self._make_api_client([self._make_task_instance("task_a")]) + + task_command.states_for_dag_run( + self.parser.parse_args( + ["tasks", "states-for-dag-run", self.dag_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.list.assert_called_once_with(dag_id=self.dag_id, dag_run_id=self.run_id) + mock_console_cls.return_value.print_as.assert_called_once_with( + data=[ + { + "dag_id": self.dag_id, + "logical_date": "2025-01-01T00:00:00+00:00", + "task_id": "task_a", + "state": "success", + "start_date": "", + "end_date": "", + }, + ], + output="json", + ) + + @mock.patch("airflowctl.ctl.commands.task_command.AirflowConsole") + def test_states_for_dag_run_includes_map_index_for_mapped_instances(self, mock_console_cls): + api_client = self._make_api_client( + [ + self._make_task_instance("plain_task"), + self._make_task_instance("mapped_task", map_index=0), + self._make_task_instance("mapped_task", map_index=1), + ] + ) + + task_command.states_for_dag_run( + self.parser.parse_args(["tasks", "states-for-dag-run", self.dag_id, self.run_id]), + api_client=api_client, + ) + + rows = mock_console_cls.return_value.print_as.call_args.kwargs["data"] + assert [row["map_index"] for row in rows] == ["", "0", "1"] + + @pytest.mark.parametrize( + "extra_args", + [ + [], + ["test_run", "--logical-date", "2025-01-01T00:00:00+00:00"], + ], + ids=["neither", "both"], + ) + def test_states_for_dag_run_requires_exactly_one_of_run_id_and_logical_date(self, extra_args, capsys): + api_client = mock.MagicMock() + + with pytest.raises(SystemExit, match="1"): + task_command.states_for_dag_run( + self.parser.parse_args(["tasks", "states-for-dag-run", self.dag_id, *extra_args]), + api_client=api_client, + ) + + api_client.task_instances.list.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_states_for_dag_run_rejects_bad_logical_date(self, logical_date, expected_message, capsys): + api_client = mock.MagicMock() + + with pytest.raises(SystemExit, match="1"): + task_command.states_for_dag_run( + self.parser.parse_args( + ["tasks", "states-for-dag-run", self.dag_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_states_for_dag_run_dag_run_not_found_by_logical_date(self, list_failure, capsys): + api_client = mock.MagicMock() + 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.states_for_dag_run( + self.parser.parse_args( + [ + "tasks", + "states-for-dag-run", + self.dag_id, + "--logical-date", + self.logical_date.isoformat(), + ] + ), + api_client=api_client, + ) + + api_client.task_instances.list.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" + ) + + def test_states_for_dag_run_dag_run_not_found(self, capsys): + api_client = mock.MagicMock() + api_client.task_instances.list.side_effect = _make_server_error(404) + + with pytest.raises(SystemExit, match="1"): + task_command.states_for_dag_run( + self.parser.parse_args(["tasks", "states-for-dag-run", self.dag_id, self.run_id]), + api_client=api_client, + ) + + assert _normalize_rich_output(capsys.readouterr().out) == ( + "Dag run 'test_run' of Dag 'test_dag' not found" + ) + + @pytest.mark.parametrize( + "failing_call", + ["dag_runs_list", "task_instances_list"], + ) + def test_states_for_dag_run_propagates_non_404_api_error(self, failing_call): + api_client = mock.MagicMock() + error = _make_server_error(500) + if failing_call == "dag_runs_list": + api_client.dag_runs.list.side_effect = error + argv = ["tasks", "states-for-dag-run", self.dag_id, "--logical-date", "2025-01-01T00:00:00+00:00"] + else: + api_client.task_instances.list.side_effect = error + argv = ["tasks", "states-for-dag-run", self.dag_id, self.run_id] + + with pytest.raises(ServerResponseError) as ctx: + task_command.states_for_dag_run(self.parser.parse_args(argv), api_client=api_client) + + assert ctx.value is error 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..b843b3b75682f 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,7 @@ 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", "list", "List all task instances for a given Dag run"), ], ) def test_help_texts_used_for_auto_generated_commands(self, group_name, subcommand_name, expected_help): @@ -797,3 +798,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",