From 7c72ce17e83b0411ab202d5b5fc1b85f59cb7243 Mon Sep 17 00:00:00 2001 From: haseebmalik18 Date: Wed, 5 Aug 2026 17:17:01 -0400 Subject: [PATCH] Add airflowctl tasks state command --- .../test_airflowctl_commands.py | 2 + airflow-ctl/docs/images/command_hashes.txt | 2 +- airflow-ctl/docs/images/output_tasks.svg | 78 +++--- airflow-ctl/src/airflowctl/ctl/cli_config.py | 16 ++ .../airflowctl/ctl/commands/task_command.py | 41 ++- .../ctl/commands/test_task_command.py | 239 ++++++++++++++++++ 6 files changed, 337 insertions(+), 41 deletions(-) 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 14a500c4d5bf7..2cef080dd66af 100644 --- a/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py +++ b/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py @@ -101,6 +101,8 @@ def 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 state example_bash_operator runme_0 "manual__{date_param}"', + 'tasks state example_bash_operator runme_0 --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', diff --git a/airflow-ctl/docs/images/command_hashes.txt b/airflow-ctl/docs/images/command_hashes.txt index 65c22154bc3c0..313629af4a626 100644 --- a/airflow-ctl/docs/images/command_hashes.txt +++ b/airflow-ctl/docs/images/command_hashes.txt @@ -10,7 +10,7 @@ jobs:a5b644c5da8889443bb40ee10b599270 pools:19efe105b9515ab1926ebcaf0e028d71 providers:34502fe09dc0b8b0a13e7e46efdffda6 taskinstances:bea84117114c2438eb7e7026f6bb7042 -tasks:ea587dc805cadbce81cd640d357f2661 +tasks:eb70701dfe1b9baeda39d5eea07fde5f variables:f8fc76d3d398b2780f4e97f7cd816646 version:31f4efdf8de0dbaaa4fac71ff7efecc3 plugins:4864fd8f356704bd2b3cd1aec3567e35 diff --git a/airflow-ctl/docs/images/output_tasks.svg b/airflow-ctl/docs/images/output_tasks.svg index 9f30d658f385e..08d22e642ed12 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 -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 + + 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 +stateGet the state of 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/ctl/cli_config.py b/airflow-ctl/src/airflowctl/ctl/cli_config.py index 787a53ce074b6..fa238b0305864 100755 --- a/airflow-ctl/src/airflowctl/ctl/cli_config.py +++ b/airflow-ctl/src/airflowctl/ctl/cli_config.py @@ -1159,6 +1159,22 @@ def merge_commands( ARG_MAP_INDEX, ), ), + ActionCommand( + name="state", + help="Get the state of a task instance", + description=( + "Get the state of a task instance. " + "Select the run with either run_id or --logical-date (pass exactly one)." + ), + func=lazy_load_command("airflowctl.ctl.commands.task_command.state"), + 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 abd61b23dde99..b087169fd6897 100644 --- a/airflow-ctl/src/airflowctl/ctl/commands/task_command.py +++ b/airflow-ctl/src/airflowctl/ctl/commands/task_command.py @@ -61,6 +61,15 @@ def _find_run_id_by_logical_date(api_client, dag_id: str, value: str) -> str: return dag_runs[0].dag_run_id +def _task_instance_not_found_message(dag_id: str, run_id: str, task_id: str, map_index: int) -> str: + """Build the message shown when a task instance is not found.""" + map_index_part = f" with map index {map_index}" if map_index >= 0 else "" + return ( + f"Task instance for task {task_id!r}{map_index_part} in Dag run " + f"{run_id!r} of Dag {dag_id!r} not found" + ) + + def _format_task_instance(ti: TaskInstanceResponse, has_mapped_instances: bool) -> dict[str, str]: data = { "dag_id": ti.dag_id, @@ -109,10 +118,8 @@ def failed_deps(args, api_client=NEW_API_CLIENT) -> None: ) 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]" + f"[red]{_task_instance_not_found_message(args.dag_id, run_id, args.task_id, args.map_index)}[/red]" ) sys.exit(1) raise @@ -149,3 +156,31 @@ def states_for_dag_run(args, api_client=NEW_API_CLIENT) -> None: data=[_format_task_instance(ti, has_mapped_instances) for ti in task_instances], output=args.output, ) + + +@provide_api_client(kind=ClientKind.CLI) +def state(args, api_client=NEW_API_CLIENT) -> None: + """Get the state of a task instance.""" + 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_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: + rich.print( + f"[red]{_task_instance_not_found_message(args.dag_id, run_id, args.task_id, args.map_index)}[/red]" + ) + sys.exit(1) + raise + + print(task_instance.state.value if task_instance.state else None) 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 2f6f23f6120ad..f870f3f1795ac 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 @@ -631,3 +631,242 @@ def test_states_for_dag_run_propagates_non_404_api_error(self, failing_call): task_command.states_for_dag_run(self.parser.parse_args(argv), api_client=api_client) assert ctx.value is error + + +class TestState: + 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, + logical_date=self.logical_date, + run_after=self.logical_date, + state=state, + start_date=None, + end_date=None, + 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="{}", + duration=None, + hostname=None, + unixname=None, + queue=None, + priority_weight=None, + operator=None, + operator_name=None, + queued_when=None, + scheduled_when=None, + pid=None, + executor=None, + note=None, + rendered_map_index=None, + trigger=None, + triggerer_job=None, + dag_version=None, + ) + + def _make_api_client(self, state: TaskInstanceState | None = TaskInstanceState.SUCCESS) -> 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.return_value = self._make_task_instance(state=state) + return api_client + + def test_state_by_run_id(self, capsys): + api_client = self._make_api_client(state=TaskInstanceState.SUCCESS) + + task_command.state( + self.parser.parse_args(["tasks", "state", self.dag_id, self.task_id, self.run_id]), + api_client=api_client, + ) + + api_client.dag_runs.list.assert_not_called() + 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 == "success\n" + + def test_state_by_logical_date(self, capsys): + api_client = self._make_api_client(state=TaskInstanceState.RUNNING) + + task_command.state( + self.parser.parse_args( + [ + "tasks", + "state", + 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.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 == "running\n" + + def test_state_with_map_index(self, capsys): + api_client = self._make_api_client(state=TaskInstanceState.SUCCESS) + + task_command.state( + self.parser.parse_args( + ["tasks", "state", self.dag_id, self.task_id, self.run_id, "--map-index", "3"] + ), + 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=3, + suppress_error_log=True, + ) + assert capsys.readouterr().out == "success\n" + + def test_state_prints_none_when_task_instance_has_no_state(self, capsys): + api_client = self._make_api_client(state=None) + + task_command.state( + self.parser.parse_args(["tasks", "state", self.dag_id, self.task_id, self.run_id]), + api_client=api_client, + ) + + assert capsys.readouterr().out == "None\n" + + @pytest.mark.parametrize( + "extra_args", + [ + [], + ["test_run", "--logical-date", "2025-01-01T00:00:00+00:00"], + ], + ids=["neither", "both"], + ) + def test_state_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.state( + self.parser.parse_args(["tasks", "state", self.dag_id, self.task_id, *extra_args]), + api_client=api_client, + ) + + api_client.task_instances.get.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_state_rejects_bad_logical_date(self, logical_date, expected_message, capsys): + api_client = self._make_api_client() + + with pytest.raises(SystemExit, match="1"): + task_command.state( + self.parser.parse_args( + ["tasks", "state", 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_state_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.state( + self.parser.parse_args( + [ + "tasks", + "state", + self.dag_id, + self.task_id, + "--logical-date", + self.logical_date.isoformat(), + ] + ), + api_client=api_client, + ) + + api_client.task_instances.get.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_state_task_instance_not_found(self, extra_args, expected_message, 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.state( + self.parser.parse_args( + ["tasks", "state", 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_state_reraises_non_404_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.state( + self.parser.parse_args(["tasks", "state", self.dag_id, self.task_id, self.run_id]), + api_client=api_client, + )