Remove no-op overrides and deduplicate Dag run lookup in airflowctl

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()

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()

airflow-ctl/src/airflowctl/ctl/commands/dag_command.py

@@ -27,17 +27,16 @@
 
 from airflowctl.api.client import (
     NEW_API_CLIENT,
-    Client,
     ClientKind,
     ServerResponseError,
     provide_api_client,
 )
 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:

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

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

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(