This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6e48e0fb2ba Deduplicate the Dag run lookup in airflowctl dag and task
commands (#70904)
6e48e0fb2ba is described below
commit 6e48e0fb2ba4743dfc478186eb17205f369c3ab0
Author: rjgoyln <[email protected]>
AuthorDate: Wed Sep 23 07:06:38 2026 +0800
Deduplicate the Dag run lookup in airflowctl dag and task commands (#70904)
The dags and tasks commands each carried their own copy of the same
logical-date Dag run lookup, differing only in what they returned, and each
repeated the run-selector guard in front of it. Sharing them lets the
selector
handling live in one place rather than 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.
---
.../src/airflowctl/ctl/commands/dag_command.py | 53 +----------
.../src/airflowctl/ctl/commands/task_command.py | 44 +--------
airflow-ctl/src/airflowctl/ctl/utils/dag_run.py | 100 +++++++++++++++++++++
.../airflow_ctl/ctl/commands/test_task_command.py | 12 +++
4 files changed, 117 insertions(+), 92 deletions(-)
diff --git a/airflow-ctl/src/airflowctl/ctl/commands/dag_command.py
b/airflow-ctl/src/airflowctl/ctl/commands/dag_command.py
index 12ba2ccc873..95770afa81d 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 rich.text import Text
from airflowctl.api.client import (
NEW_API_CLIENT,
- Client,
ClientKind,
ServerResponseError,
provide_api_client,
@@ -35,10 +34,10 @@ from airflowctl.api.client import (
from airflowctl.api.datamodels.generated import (
ClearTaskInstancesBody,
DAGPatchBody,
- DAGRunResponse,
DagSchedulingState,
)
from airflowctl.ctl.console_formatting import AirflowConsole
+from airflowctl.ctl.utils.dag_run import resolve_dag_run
def update_dag_state(
@@ -133,58 +132,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 abd61b23dde..1c6a7004903 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 @@ import rich
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 00000000000..b2bdad4af38
--- /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 2f6f23f6120..9b15695209c 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 @@ class TestFailedDeps:
"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 @@ class TestStatesForDagRun:
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(
