Handle null logical date in CLI commands

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",

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,17 @@ 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:
-        conf_out = ", " + json.dumps(dr.conf)
-    print(str(out) + conf_out)
+    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 not dr:
+        print(None)
+    elif dr.conf:
+        print(f"{dr.state}, {json.dumps(dr.conf)}")
+    else:
+        print(dr.state)
 
 
 @cli_utils.action_cli
@@ -465,20 +471,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 +521,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()
 

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,
@@ -165,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(
@@ -592,17 +552,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 +569,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 "",

airflow/cli/utils.py

@@ -17,13 +17,17 @@
 
 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 sqlalchemy.orm import Session
+
+    from airflow.models.dagrun import DagRun
 
 
 class CliConflictError(Exception):
@@ -45,8 +49,50 @@ 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:
         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.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
+    try:
+        logical_date = timezone.parse(value)
+    except (ParserError, TypeError):
+        return None, 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

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"]
         )

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