From ea8a0ce49388a01686b0639b89f6c33ef8a7bf3f Mon Sep 17 00:00:00 2001 From: rjgoyln Date: Sat, 1 Aug 2026 18:41:06 +0800 Subject: [PATCH] Remove no-op overrides and deduplicate Dag run lookup in airflowctl Three pieces of airflowctl did nothing but obscure the code around them: an except clause that re-raised its exception unchanged, an argparse override that only called super(), and a Dag run lookup that existed in two commands. The re-raise is the same pattern cleaned out of the operations layer in #70376; it survived there because it sits in client.py instead. Sharing the lookup also lets the run selector validation live in one place rather than being repeated in front of every command that accepts one. A supplied run_id is deliberately still taken at face value rather than fetched, so commands acting on a nested resource keep reporting a miss against that resource instead of against the Dag run. The added assertions pin that down, since collapsing the two resolvers would otherwise silently add a request and re-attribute the 404. --- airflow-ctl/src/airflowctl/api/client.py | 3 - airflow-ctl/src/airflowctl/ctl/cli_config.py | 4 - .../airflowctl/ctl/commands/dag_command.py | 53 +--------- .../airflowctl/ctl/commands/task_command.py | 44 +------- .../src/airflowctl/ctl/utils/dag_run.py | 100 ++++++++++++++++++ .../ctl/commands/test_task_command.py | 12 +++ 6 files changed, 117 insertions(+), 99 deletions(-) create mode 100644 airflow-ctl/src/airflowctl/ctl/utils/dag_run.py diff --git a/airflow-ctl/src/airflowctl/api/client.py b/airflow-ctl/src/airflowctl/api/client.py index 65f90302868e5..51e92d813750f 100644 --- a/airflow-ctl/src/airflowctl/api/client.py +++ b/airflow-ctl/src/airflowctl/api/client.py @@ -67,7 +67,6 @@ AirflowCtlCredentialNotFoundException, AirflowCtlException, AirflowCtlKeyringException, - AirflowCtlNotFoundException, ) if TYPE_CHECKING: @@ -517,8 +516,6 @@ def get_client( kind=kind, ) yield api_client - except AirflowCtlNotFoundException as e: - raise e finally: if api_client: api_client.close() diff --git a/airflow-ctl/src/airflowctl/ctl/cli_config.py b/airflow-ctl/src/airflowctl/ctl/cli_config.py index c3057099dd875..787a53ce074b6 100755 --- a/airflow-ctl/src/airflowctl/ctl/cli_config.py +++ b/airflow-ctl/src/airflowctl/ctl/cli_config.py @@ -111,10 +111,6 @@ def safe_call_command(function: Callable, args: Iterable[Arg]) -> None: class DefaultHelpParser(argparse.ArgumentParser): """CustomParser to display help message.""" - def _check_value(self, action, value): - """Override _check_value and check conditionally added command.""" - super()._check_value(action, value) - def error(self, message): """Override error and use print_help instead of print_usage.""" self.print_help() diff --git a/airflow-ctl/src/airflowctl/ctl/commands/dag_command.py b/airflow-ctl/src/airflowctl/ctl/commands/dag_command.py index 15307085f3612..4bcb764e66493 100644 --- a/airflow-ctl/src/airflowctl/ctl/commands/dag_command.py +++ b/airflow-ctl/src/airflowctl/ctl/commands/dag_command.py @@ -27,7 +27,6 @@ from airflowctl.api.client import ( NEW_API_CLIENT, - Client, ClientKind, ServerResponseError, provide_api_client, @@ -35,9 +34,9 @@ from airflowctl.api.datamodels.generated import ( ClearTaskInstancesBody, DAGPatchBody, - DAGRunResponse, ) from airflowctl.ctl.console_formatting import AirflowConsole +from airflowctl.ctl.utils.dag_run import resolve_dag_run def update_dag_state( @@ -118,58 +117,10 @@ def next_execution(args, api_client=NEW_API_CLIENT) -> dict | None: return result -def _get_dag_run_by_run_id(api_client: Client, dag_id: str, run_id: str) -> DAGRunResponse: - """Get a Dag run by its run ID.""" - try: - return api_client.dag_runs.get(dag_id=dag_id, dag_run_id=run_id, suppress_error_log=True) - except ServerResponseError as e: - if e.response.status_code != 404: - raise - rich.print(f"[red]Dag run {run_id!r} of Dag {dag_id!r} not found[/red]") - sys.exit(1) - - -def _get_dag_run_by_logical_date(api_client: Client, dag_id: str, value: str) -> DAGRunResponse: - """Get 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] - - @provide_api_client(kind=ClientKind.CLI) def state(args, api_client=NEW_API_CLIENT) -> None: """Show the state and configuration of 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) - - if args.run_id: - dag_run = _get_dag_run_by_run_id(api_client, args.dag_id, args.run_id) - else: - dag_run = _get_dag_run_by_logical_date(api_client, args.dag_id, args.logical_date) + dag_run = resolve_dag_run(api_client, args) state_value = getattr(dag_run.state, "value", dag_run.state) if dag_run.conf: diff --git a/airflow-ctl/src/airflowctl/ctl/commands/task_command.py b/airflow-ctl/src/airflowctl/ctl/commands/task_command.py index abd61b23dde99..1c6a700490320 100644 --- a/airflow-ctl/src/airflowctl/ctl/commands/task_command.py +++ b/airflow-ctl/src/airflowctl/ctl/commands/task_command.py @@ -17,7 +17,6 @@ from __future__ import annotations -import datetime import sys from typing import TYPE_CHECKING @@ -26,41 +25,12 @@ 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 +from airflowctl.ctl.utils.dag_run import resolve_dag_run_id 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, @@ -78,11 +48,7 @@ def _format_task_instance(ti: TaskInstanceResponse, has_mapped_instances: bool) @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) + run_id = resolve_dag_run_id(api_client, args) try: response = api_client.task_instances.get_dependencies( @@ -129,11 +95,7 @@ def failed_deps(args, api_client=NEW_API_CLIENT) -> None: @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) + run_id = resolve_dag_run_id(api_client, args) try: task_instances = api_client.task_instances.list(dag_id=args.dag_id, dag_run_id=run_id).task_instances diff --git a/airflow-ctl/src/airflowctl/ctl/utils/dag_run.py b/airflow-ctl/src/airflowctl/ctl/utils/dag_run.py new file mode 100644 index 0000000000000..b2bdad4af3809 --- /dev/null +++ b/airflow-ctl/src/airflowctl/ctl/utils/dag_run.py @@ -0,0 +1,100 @@ +# 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. +"""Resolve the Dag run a command acts on from its ``run_id`` / ``--logical-date`` selectors.""" + +from __future__ import annotations + +import datetime +import sys +from typing import TYPE_CHECKING + +import rich + +from airflowctl.api.client import ServerResponseError + +if TYPE_CHECKING: + from airflowctl.api.client import Client + from airflowctl.api.datamodels.generated import DAGRunResponse + + +def _validate_selectors(args) -> None: + """Exit unless exactly one of ``run_id`` and ``--logical-date`` was given.""" + 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) + + +def _get_dag_run_by_run_id(api_client: Client, dag_id: str, run_id: str) -> DAGRunResponse: + """Get a Dag run by its run ID.""" + try: + return api_client.dag_runs.get(dag_id=dag_id, dag_run_id=run_id, suppress_error_log=True) + except ServerResponseError as e: + if e.response.status_code != 404: + raise + rich.print(f"[red]Dag run {run_id!r} of Dag {dag_id!r} not found[/red]") + sys.exit(1) + + +def _get_dag_run_by_logical_date(api_client: Client, dag_id: str, value: str) -> DAGRunResponse: + """Get 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] + + +def resolve_dag_run(api_client: Client, args) -> DAGRunResponse: + """Get the selected Dag run, fetching it when ``run_id`` was given.""" + _validate_selectors(args) + if args.run_id: + return _get_dag_run_by_run_id(api_client, args.dag_id, args.run_id) + return _get_dag_run_by_logical_date(api_client, args.dag_id, args.logical_date) + + +def resolve_dag_run_id(api_client: Client, args) -> str: + """ + Get the ID of the selected Dag run. + + A ``run_id`` is taken at face value rather than fetched, so that a caller acting on a nested + resource reports the miss against that resource instead of against the Dag run. + """ + _validate_selectors(args) + if args.run_id: + return args.run_id + return _get_dag_run_by_logical_date(api_client, args.dag_id, args.logical_date).dag_run_id 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..9b15695209c78 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 @@ -187,6 +187,17 @@ def test_failed_deps_with_map_index(self, capsys): "Task instance dependencies not met:\nPool Slots Available: pool is full\n" ) + def test_failed_deps_does_not_look_up_the_dag_run_for_a_given_run_id(self): + api_client = self._make_api_client() + + 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.dag_runs.get.assert_not_called() + api_client.dag_runs.list.assert_not_called() + 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")] @@ -452,6 +463,7 @@ def test_states_for_dag_run_by_run_id(self, mock_console_cls): api_client=api_client, ) + api_client.dag_runs.get.assert_not_called() 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(