From 84b370fe05f18253f06f3c0f1013586ed4237591 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 4 Feb 2025 12:43:01 +0800 Subject: [PATCH 1/4] Handle null logical date in CLI commands Relatively simple things. Sorting by logical date is replaced by run_after instead. Selecting a run using logical date is replaced by the smart logical-date-or-run-ID argument already available for other purposes. --- airflow/cli/cli_config.py | 2 +- .../commands/remote_commands/dag_command.py | 33 ++++++----- .../commands/remote_commands/task_command.py | 56 +++---------------- airflow/cli/utils.py | 47 ++++++++++++++++ 4 files changed, 74 insertions(+), 64 deletions(-) diff --git a/airflow/cli/cli_config.py b/airflow/cli/cli_config.py index 0186add2bb3bb..78cfbf394b839 100644 --- a/airflow/cli/cli_config.py +++ b/airflow/cli/cli_config.py @@ -1073,7 +1073,7 @@ class GroupCommand(NamedTuple): name="state", help="Get the status of a dag run", func=lazy_load_command("airflow.cli.commands.remote_commands.dag_command.dag_state"), - args=(ARG_DAG_ID, ARG_LOGICAL_DATE, ARG_SUBDIR, ARG_VERBOSE), + args=(ARG_DAG_ID, ARG_LOGICAL_DATE_OR_RUN_ID, ARG_SUBDIR, ARG_VERBOSE), ), ActionCommand( name="next-execution", diff --git a/airflow/cli/commands/remote_commands/dag_command.py b/airflow/cli/commands/remote_commands/dag_command.py index 6c95451271fbc..7a2ff67026ee0 100644 --- a/airflow/cli/commands/remote_commands/dag_command.py +++ b/airflow/cli/commands/remote_commands/dag_command.py @@ -33,6 +33,7 @@ from airflow.api.client import get_current_api_client from airflow.api_connexion.schemas.dag_schema import dag_schema from airflow.cli.simple_table import AirflowConsole +from airflow.cli.utils import fetch_dag_run_from_run_id_or_logical_date_string from airflow.exceptions import AirflowException from airflow.jobs.job import Job from airflow.models import DagBag, DagModel, DagRun, TaskInstance @@ -264,12 +265,16 @@ def dag_state(args, session: Session = NEW_SESSION) -> None: if not dag: raise SystemExit(f"DAG: {args.dag_id} does not exist in 'dag' table") - dr = session.scalar(select(DagRun).filter_by(dag_id=args.dag_id, logical_date=args.logical_date)) - out = dr.state if dr else None - conf_out = "" - if out and dr.conf: + dr, _ = fetch_dag_run_from_run_id_or_logical_date_string( + dag_id=dag.dag_id, + value=args.logical_date_or_run_id, + session=session, + ) + if dr and dr.conf: conf_out = ", " + json.dumps(dr.conf) - print(str(out) + conf_out) + else: + conf_out = "" + print(f"{dr.state if dr else None}{conf_out}") @cli_utils.action_cli @@ -465,20 +470,20 @@ def dag_list_dag_runs(args, dag: DAG | None = None, session: Session = NEW_SESSI logical_end_date=args.end_date, session=session, ) + dag_runs.sort(key=operator.attrgetter("run_after"), reverse=True) - dag_runs.sort(key=lambda x: x.logical_date, reverse=True) - AirflowConsole().print_as( - data=dag_runs, - output=args.output, - mapper=lambda dr: { + def _render_dagrun(dr: DagRun) -> dict[str, str]: + return { "dag_id": dr.dag_id, "run_id": dr.run_id, "state": dr.state, - "logical_date": dr.logical_date.isoformat(), + "run_after": dr.run_after.isoformat(), + "logical_date": dr.logical_date.isoformat() if dr.logical_date else "", "start_date": dr.start_date.isoformat() if dr.start_date else "", "end_date": dr.end_date.isoformat() if dr.end_date else "", - }, - ) + } + + AirflowConsole().print_as(data=dag_runs, output=args.output, mapper=_render_dagrun) @cli_utils.action_cli @@ -515,7 +520,7 @@ def dag_test(args, dag: DAG | None = None, session: Session = NEW_SESSION) -> No tis = session.scalars( select(TaskInstance).where( TaskInstance.dag_id == args.dag_id, - TaskInstance.logical_date == logical_date, + TaskInstance.run_id == dr.run_id, ) ).all() diff --git a/airflow/cli/commands/remote_commands/task_command.py b/airflow/cli/commands/remote_commands/task_command.py index ccae29ededd09..58dab162705c2 100644 --- a/airflow/cli/commands/remote_commands/task_command.py +++ b/airflow/cli/commands/remote_commands/task_command.py @@ -32,11 +32,10 @@ from typing import TYPE_CHECKING, Protocol, cast import pendulum -from pendulum.parsing.exceptions import ParserError -from sqlalchemy import select from airflow import settings from airflow.cli.simple_table import AirflowConsole +from airflow.cli.utils import fetch_dag_run_from_run_id_or_logical_date_string from airflow.configuration import conf from airflow.exceptions import AirflowException, DagRunNotFound, TaskDeferred, TaskInstanceNotFound from airflow.executors.executor_loader import ExecutorLoader @@ -91,41 +90,6 @@ def _generate_temporary_run_id() -> str: return f"__airflow_temporary_run_{timezone.utcnow().isoformat()}__" -def _fetch_dag_run_from_run_id_or_logical_date_string( - *, - dag_id: str, - value: str, - session: Session, -) -> tuple[DagRun, pendulum.DateTime | None]: - """ - Try to find a DAG run with a given string value. - - The string value may be a run ID, or a logical date in string form. We first - try to use it as a run_id; if a run is found, it is returned as-is. - - Otherwise, the string value is parsed into a datetime. If that works, it is - used to find a DAG run. - - The return value is a two-tuple. The first item is the found DAG run (or - *None* if one cannot be found). The second is the parsed logical date. This - second value can be used to create a new run by the calling function when - one cannot be found here. - """ - if dag_run := DAG.fetch_dagrun(dag_id=dag_id, run_id=value, session=session): - return dag_run, dag_run.logical_date # type: ignore[return-value] - try: - logical_date = timezone.parse(value) - except (ParserError, TypeError): - return dag_run, None - dag_run = session.scalar( - select(DagRun) - .where(DagRun.dag_id == dag_id, DagRun.logical_date == logical_date) - .order_by(DagRun.id.desc()) - .limit(1) - ) - return dag_run, logical_date - - def _get_dag_run( *, dag: DAG, @@ -152,7 +116,7 @@ def _get_dag_run( logical_date = None if logical_date_or_run_id: - dag_run, logical_date = _fetch_dag_run_from_run_id_or_logical_date_string( + dag_run, logical_date = fetch_dag_run_from_run_id_or_logical_date_string( dag_id=dag.dag_id, value=logical_date_or_run_id, session=session, @@ -592,17 +556,11 @@ def _guess_debugger() -> _SupportedDebugger: @provide_session def task_states_for_dag_run(args, session: Session = NEW_SESSION) -> None: """Get the status of all task instances in a DagRun.""" - dag_run = session.scalar( - select(DagRun).where(DagRun.run_id == args.logical_date_or_run_id, DagRun.dag_id == args.dag_id) + dag_run, _ = fetch_dag_run_from_run_id_or_logical_date_string( + dag_id=args.dag_id, + value=args.logical_date_or_run_id, + session=session, ) - if not dag_run: - try: - logical_date = timezone.parse(args.logical_date_or_run_id) - dag_run = session.scalar( - select(DagRun).where(DagRun.logical_date == logical_date, DagRun.dag_id == args.dag_id) - ) - except (ParserError, TypeError) as err: - raise AirflowException(f"Error parsing the supplied logical_date. Error: {err}") if dag_run is None: raise DagRunNotFound( @@ -615,7 +573,7 @@ def task_states_for_dag_run(args, session: Session = NEW_SESSION) -> None: def format_task_instance(ti: TaskInstance) -> dict[str, str]: data = { "dag_id": ti.dag_id, - "logical_date": dag_run.logical_date.isoformat(), + "logical_date": dag_run.logical_date.isoformat() if dag_run.logical_date else "", "task_id": ti.task_id, "state": ti.state, "start_date": ti.start_date.isoformat() if ti.start_date else "", diff --git a/airflow/cli/utils.py b/airflow/cli/utils.py index d132deeb373a9..04c2e42f38cb9 100644 --- a/airflow/cli/utils.py +++ b/airflow/cli/utils.py @@ -25,6 +25,11 @@ if TYPE_CHECKING: from io import IOBase + from pendulum import DateTime + from sqlalchemy.orm import Session + + from airflow.models.dagrun import DagRun + class CliConflictError(Exception): """Error for when CLI commands are defined twice by different sources.""" @@ -50,3 +55,45 @@ def print_export_output(command_type: str, exported_items: Collection, file: io. print(f"\n{len(exported_items)} {command_type} successfully exported.", file=sys.stderr) else: print(f"{len(exported_items)} {command_type} successfully exported to {file.name}.") + + +def fetch_dag_run_from_run_id_or_logical_date_string( + *, + dag_id: str, + value: str, + session: Session, +) -> tuple[DagRun | None, DateTime | None]: + """ + Try to find a DAG run with a given string value. + + The string value may be a run ID, or a logical date in string form. We first + try to use it as a run_id; if a run is found, it is returned as-is. + + Otherwise, the string value is parsed into a datetime. If that works, it is + used to find a DAG run. + + The return value is a two-tuple. The first item is the found DAG run (or + *None* if one cannot be found). The second is the parsed logical date. This + second value can be used to create a new run by the calling function when + one cannot be found here. + """ + from pendulum.parsing.exceptions import ParserError + from sqlalchemy import select + + from airflow.models.dag import DAG + from airflow.models.dagrun import DagRun + from airflow.utils import timezone + + if dag_run := DAG.fetch_dagrun(dag_id=dag_id, run_id=value, session=session): + return dag_run, dag_run.logical_date # type: ignore[return-value] + try: + logical_date = timezone.parse(value) + except (ParserError, TypeError): + return dag_run, None + dag_run = session.scalar( + select(DagRun) + .where(DagRun.dag_id == dag_id, DagRun.logical_date == logical_date) + .order_by(DagRun.id.desc()) + .limit(1) + ) + return dag_run, logical_date From 5fb571ba5482f52befba77bef5bdd734bb091c72 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 4 Feb 2025 15:47:05 +0800 Subject: [PATCH 2/4] Fix mock setups in tests --- .../remote_commands/test_dag_command.py | 2 ++ .../remote_commands/test_task_command.py | 20 ++++++------------- 2 files changed, 8 insertions(+), 14 deletions(-) diff --git a/tests/cli/commands/remote_commands/test_dag_command.py b/tests/cli/commands/remote_commands/test_dag_command.py index 04b0ef88104e7..addae8299c3e7 100644 --- a/tests/cli/commands/remote_commands/test_dag_command.py +++ b/tests/cli/commands/remote_commands/test_dag_command.py @@ -693,6 +693,8 @@ def test_dag_test_conf(self, mock_get_dag): ) @mock.patch("airflow.cli.commands.remote_commands.dag_command.get_dag") def test_dag_test_show_dag(self, mock_get_dag, mock_render_dag): + mock_get_dag.return_value.test.return_value.run_id = "__test_dag_test_show_dag_fake_dag_run_run_id__" + cli_args = self.parser.parse_args( ["dags", "test", "example_bash_operator", DEFAULT_DATE.isoformat(), "--show-dagrun"] ) diff --git a/tests/cli/commands/remote_commands/test_task_command.py b/tests/cli/commands/remote_commands/test_task_command.py index 166f152b9eaee..98e31966b904d 100644 --- a/tests/cli/commands/remote_commands/test_task_command.py +++ b/tests/cli/commands/remote_commands/test_task_command.py @@ -228,30 +228,22 @@ def test_cli_test_different_path(self, session, tmp_path): # verify that the file was in different location when run assert ti.xcom_pull(ti.task_id) == new_file_path.as_posix() - @mock.patch("airflow.cli.commands.remote_commands.task_command.select") - @mock.patch("sqlalchemy.orm.session.Session.scalar") - def test_task_render_with_custom_timetable(self, mock_scalar, mock_select): + @mock.patch( + "airflow.cli.commands.remote_commands.task_command.fetch_dag_run_from_run_id_or_logical_date_string" + ) + def test_task_render_with_custom_timetable(self, mock_fetch_dag_run_from_run_id_or_logical_date_string): """ Test that the `tasks render` CLI command queries the database correctly for a DAG with a custom timetable. Verifies that a query is executed to fetch the appropriate DagRun and that the database interaction occurs as expected. """ - from sqlalchemy import select - - from airflow.models.dagrun import DagRun - - mock_query = ( - select(DagRun).where(DagRun.dag_id == "example_workday_timetable").order_by(DagRun.id.desc()) - ) - mock_select.return_value = mock_query - - mock_scalar.return_value = None + mock_fetch_dag_run_from_run_id_or_logical_date_string.return_value = (None, None) task_command.task_render( self.parser.parse_args(["tasks", "render", "example_workday_timetable", "run_this", "2022-01-01"]) ) - mock_select.assert_called_once() + mock_fetch_dag_run_from_run_id_or_logical_date_string.assert_called_once() @pytest.mark.filterwarnings("ignore::airflow.utils.context.AirflowContextDeprecationWarning") def test_test_with_existing_dag_run(self, caplog): From 899e55ac3331639cc3c9183098b03b65751eb046 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Thu, 6 Feb 2025 16:26:34 +0800 Subject: [PATCH 3/4] Improve typing --- .../cli/commands/remote_commands/task_command.py | 6 +----- airflow/cli/utils.py | 15 +++++++-------- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/airflow/cli/commands/remote_commands/task_command.py b/airflow/cli/commands/remote_commands/task_command.py index 58dab162705c2..296bba5965ccd 100644 --- a/airflow/cli/commands/remote_commands/task_command.py +++ b/airflow/cli/commands/remote_commands/task_command.py @@ -129,11 +129,7 @@ def _get_dag_run( f"of {logical_date_or_run_id!r} not found" ) - if logical_date is not None: - dag_run_logical_date = logical_date - else: - dag_run_logical_date = pendulum.instance(timezone.utcnow()) - + dag_run_logical_date = pendulum.instance(logical_date or timezone.utcnow()) if create_if_necessary == "memory": data_interval = dag.timetable.infer_manual_data_interval(run_after=dag_run_logical_date) dag_run = DagRun( diff --git a/airflow/cli/utils.py b/airflow/cli/utils.py index 04c2e42f38cb9..605244ee71e56 100644 --- a/airflow/cli/utils.py +++ b/airflow/cli/utils.py @@ -17,15 +17,14 @@ from __future__ import annotations -import io import sys -from collections.abc import Collection from typing import TYPE_CHECKING if TYPE_CHECKING: - from io import IOBase + import datetime + from collections.abc import Collection + from io import IOBase, TextIOWrapper - from pendulum import DateTime from sqlalchemy.orm import Session from airflow.models.dagrun import DagRun @@ -50,7 +49,7 @@ def is_stdout(fileio: IOBase) -> bool: return fileio.fileno() == sys.stdout.fileno() -def print_export_output(command_type: str, exported_items: Collection, file: io.TextIOWrapper): +def print_export_output(command_type: str, exported_items: Collection, file: TextIOWrapper): if not file.closed and is_stdout(file): print(f"\n{len(exported_items)} {command_type} successfully exported.", file=sys.stderr) else: @@ -62,7 +61,7 @@ def fetch_dag_run_from_run_id_or_logical_date_string( dag_id: str, value: str, session: Session, -) -> tuple[DagRun | None, DateTime | None]: +) -> tuple[DagRun | None, datetime.datetime | None]: """ Try to find a DAG run with a given string value. @@ -85,11 +84,11 @@ def fetch_dag_run_from_run_id_or_logical_date_string( from airflow.utils import timezone if dag_run := DAG.fetch_dagrun(dag_id=dag_id, run_id=value, session=session): - return dag_run, dag_run.logical_date # type: ignore[return-value] + return dag_run, dag_run.logical_date try: logical_date = timezone.parse(value) except (ParserError, TypeError): - return dag_run, None + return None, None dag_run = session.scalar( select(DagRun) .where(DagRun.dag_id == dag_id, DagRun.logical_date == logical_date) From 7b32dfa46292feb4a20268bf13ccbb264c088964 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Thu, 6 Feb 2025 16:29:40 +0800 Subject: [PATCH 4/4] Better string formation --- airflow/cli/commands/remote_commands/dag_command.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/airflow/cli/commands/remote_commands/dag_command.py b/airflow/cli/commands/remote_commands/dag_command.py index 7a2ff67026ee0..bf0841f1250be 100644 --- a/airflow/cli/commands/remote_commands/dag_command.py +++ b/airflow/cli/commands/remote_commands/dag_command.py @@ -270,11 +270,12 @@ def dag_state(args, session: Session = NEW_SESSION) -> None: value=args.logical_date_or_run_id, session=session, ) - if dr and dr.conf: - conf_out = ", " + json.dumps(dr.conf) + if not dr: + print(None) + elif dr.conf: + print(f"{dr.state}, {json.dumps(dr.conf)}") else: - conf_out = "" - print(f"{dr.state if dr else None}{conf_out}") + print(dr.state) @cli_utils.action_cli