Split airflowctl dags state argument into run_id and --logical-date

airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py

@@ -85,7 +85,7 @@ def date_param():
     # Order of trigger, state, and pause/unpause is important for test stability
     "dags trigger example_bash_operator --logical-date={date_param} --run-after={date_param}",
     'dags state example_bash_operator "manual__{date_param}"',
-    'dags state example_bash_operator "{date_param}"',
+    'dags state example_bash_operator --logical-date "{date_param}"',
     # Test trigger without logical-date (should default to now)
     "dags trigger example_bash_operator",
     "dags next-execution example_bash_operator",

airflow-ctl/src/airflowctl/ctl/cli_config.py

@@ -281,10 +281,16 @@ def _load_help_texts_yaml() -> dict[str, dict[str, str]]:
     type=str,
     help="The Dag ID",
 )
-ARG_LOGICAL_DATE_OR_RUN_ID = Arg(
-    flags=("logical_date_or_run_id",),
+ARG_RUN_ID = Arg(
+    flags=("run_id",),
     type=str,
-    help="The logical date with a timezone offset or run ID of the Dag run",
+    nargs="?",
+    help="The run ID of the Dag run (pass this or --logical-date, not both)",
+)
+ARG_LOGICAL_DATE = Arg(
+    flags=("--logical-date",),
+    type=str,
+    help="The logical date of the Dag run with a timezone offset (pass this or run_id, not both)",
 )
 
 # Task Commands Args
@@ -1015,7 +1021,8 @@ def merge_commands(
         func=lazy_load_command("airflowctl.ctl.commands.dag_command.state"),
         args=(
             ARG_DAG_ID,
-            ARG_LOGICAL_DATE_OR_RUN_ID,
+            ARG_RUN_ID,
+            ARG_LOGICAL_DATE,
         ),
     ),
     ActionCommand(

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

@@ -25,7 +25,13 @@
 import rich
 from rich.text import Text
 
-from airflowctl.api.client import NEW_API_CLIENT, ClientKind, ServerResponseError, provide_api_client
+from airflowctl.api.client import (
+    NEW_API_CLIENT,
+    Client,
+    ClientKind,
+    ServerResponseError,
+    provide_api_client,
+)
 from airflowctl.api.datamodels.generated import DAGPatchBody, DAGRunResponse
 from airflowctl.ctl.console_formatting import AirflowConsole
 
@@ -108,53 +114,61 @@ def next_execution(args, api_client=NEW_API_CLIENT) -> dict | None:
     return result
 
 
-def _parse_logical_date(value: str) -> datetime.datetime | None:
-    """Parse an ISO-formatted logical date."""
+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:
-        return None
+        rich.print(f"[red]Invalid --logical-date: {value!r}[/red]")
+        sys.exit(1)
     if logical_date.tzinfo is None:
-        raise SystemExit("Logical date must include a timezone offset")
-    return logical_date
-
+        rich.print("[red]--logical-date must include a timezone offset[/red]")
+        sys.exit(1)
 
-def _get_dag_run_by_run_id_or_logical_date(api_client, dag_id: str, value: str) -> DAGRunResponse | None:
-    """Get a Dag run by run ID, falling back to an exact logical date match."""
+    dag_runs = []
     try:
-        return api_client.dag_runs.get(dag_id=dag_id, dag_run_id=value, suppress_error_log=True)
-    except ServerResponseError as e:
-        if e.response.status_code != 404:
-            raise
-
-    if logical_date := _parse_logical_date(value):
-        response = api_client.dag_runs.list(
+        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,
-        )
-        if response.dag_runs:
-            return response.dag_runs[0]
-    else:
-        api_client.dag_runs.list(dag_id=dag_id, limit=1)
-    return None
+            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."""
-    dag_run = _get_dag_run_by_run_id_or_logical_date(
-        api_client=api_client,
-        dag_id=args.dag_id,
-        value=args.logical_date_or_run_id,
-    )
-    if not dag_run:
-        rich.print("[yellow]No matching Dag run found.[/yellow]")
+    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)
+
+    state_value = getattr(dag_run.state, "value", dag_run.state)
+    if dag_run.conf:
+        rich.print(Text(f"{state_value}, {json.dumps(dag_run.conf)}"))
     else:
-        state_value = getattr(dag_run.state, "value", dag_run.state)
-        if dag_run.conf:
-            rich.print(Text(f"{state_value}, {json.dumps(dag_run.conf)}"))
-        else:
-            rich.print(Text(state_value))
+        rich.print(Text(state_value))

airflow-ctl/tests/airflow_ctl/api/test_operations.py

@@ -1233,6 +1233,20 @@ def handle_request(request: httpx.Request) -> httpx.Response:
         )
         assert response == self.dag_run_collection_response
 
+    @pytest.mark.parametrize("suppress_error_log", [False, True])
+    def test_list_passes_error_log_suppression_extension(self, suppress_error_log):
+        def handle_request(request: httpx.Request) -> httpx.Response:
+            assert request.extensions["airflowctl_suppress_error_log"] is suppress_error_log
+            return httpx.Response(200, json=json.loads(self.dag_run_collection_response.model_dump_json()))
+
+        client = make_api_client(transport=httpx.MockTransport(handle_request))
+        response = client.dag_runs.list(
+            dag_id=self.dag_id,
+            limit=1,
+            suppress_error_log=suppress_error_log,
+        )
+        assert response == self.dag_run_collection_response
+
     def test_list_with_logical_date_filters_and_order(self):
         logical_date = datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc)
 

airflow-ctl/tests/airflow_ctl/ctl/commands/test_dag_command.py

@@ -29,12 +29,16 @@
 from airflowctl.ctl.commands import dag_command
 
 
-def _server_error(status_code: int) -> ServerResponseError:
+def _make_server_error(status_code: int) -> ServerResponseError:
     request = httpx.Request("GET", "http://testserver/api/v2/dags/test_dag/dagRuns/test_run")
     response = httpx.Response(status_code, request=request, json={"detail": "boom"})
     return ServerResponseError(message="boom", request=request, response=response)
 
 
+def _normalize_rich_output(text: str) -> str:
+    return " ".join(text.split())
+
+
 class TestDagCommands:
     parser = cli_parser.get_parser()
     dag_id = "test_dag"
@@ -245,13 +249,14 @@ def test_state_by_run_id(self, capsys):
     def test_state_by_logical_date(self, capsys):
         api_client = mock.MagicMock()
         logical_date = datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc)
-        api_client.dag_runs.get.side_effect = _server_error(404)
         api_client.dag_runs.list.return_value.dag_runs = [
             mock.MagicMock(state="failed", conf={"reason": "[red]test[/red]"})
         ]
 
         dag_command.state(
-            self.parser.parse_args(["dags", "state", self.dag_id, logical_date.isoformat()]),
+            self.parser.parse_args(
+                ["dags", "state", self.dag_id, "--logical-date", logical_date.isoformat()]
+            ),
             api_client=api_client,
         )
 
@@ -262,67 +267,95 @@ def test_state_by_logical_date(self, capsys):
             logical_date_lte=logical_date,
             order_by="-id",
             limit=1,
+            suppress_error_log=True,
         )
+        api_client.dag_runs.get.assert_not_called()
 
     @pytest.mark.parametrize(
-        ("value", "expected_list_kwargs"),
+        "extra_args",
         [
-            pytest.param("missing_run", {"dag_id": dag_id, "limit": 1}, id="run-id"),
-            pytest.param(
-                "2025-01-01T00:00:00+00:00",
-                {
-                    "dag_id": dag_id,
-                    "logical_date_gte": datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
-                    "logical_date_lte": datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
-                    "order_by": "-id",
-                    "limit": 1,
-                },
-                id="logical-date",
-            ),
+            [],
+            ["test_run", "--logical-date", "2025-01-01T00:00:00+00:00"],
         ],
+        ids=["neither", "both"],
     )
-    @mock.patch("rich.print")
-    def test_state_missing_run_prints_message(self, mock_rich_print, value, expected_list_kwargs):
+    def test_state_requires_exactly_one_of_run_id_and_logical_date(self, extra_args, capsys):
         api_client = mock.MagicMock()
-        api_client.dag_runs.get.side_effect = _server_error(404)
-        api_client.dag_runs.list.return_value.dag_runs = []
 
-        dag_command.state(
-            self.parser.parse_args(["dags", "state", self.dag_id, value]),
-            api_client=api_client,
+        with pytest.raises(SystemExit) as ctx:
+            dag_command.state(
+                self.parser.parse_args(["dags", "state", self.dag_id, *extra_args]),
+                api_client=api_client,
+            )
+
+        assert ctx.value.code == 1
+        api_client.dag_runs.get.assert_not_called()
+        assert _normalize_rich_output(capsys.readouterr().out) == (
+            "Provide either run_id or --logical-date, but not both"
         )
 
-        mock_rich_print.assert_called_once_with("[yellow]No matching Dag run found.[/yellow]")
-        api_client.dag_runs.list.assert_called_once_with(**expected_list_kwargs)
+    @pytest.mark.parametrize(
+        ("logical_date", "expected_message"),
+        [
+            ("not-a-date", "Invalid --logical-date: 'not-a-date'"),
+            ("2025-01-01T00:00:00", "--logical-date must include a timezone offset"),
+        ],
+        ids=["unparsable", "naive"],
+    )
+    def test_state_rejects_bad_logical_date(self, logical_date, expected_message, capsys):
+        api_client = mock.MagicMock()
 
-    def test_state_missing_dag_propagates_api_error(self):
+        with pytest.raises(SystemExit) as ctx:
+            dag_command.state(
+                self.parser.parse_args(["dags", "state", self.dag_id, "--logical-date", logical_date]),
+                api_client=api_client,
+            )
+
+        assert ctx.value.code == 1
+        api_client.dag_runs.list.assert_not_called()
+        assert _normalize_rich_output(capsys.readouterr().out) == expected_message
+
+    @pytest.mark.parametrize("list_failure", ["no_matching_run", "dag_not_found_404"])
+    def test_state_dag_run_not_found_by_logical_date(self, list_failure, capsys):
         api_client = mock.MagicMock()
-        api_client.dag_runs.get.side_effect = _server_error(404)
-        api_client.dag_runs.list.side_effect = error = _server_error(404)
+        if list_failure == "no_matching_run":
+            api_client.dag_runs.list.return_value.dag_runs = []
+        else:
+            api_client.dag_runs.list.side_effect = _make_server_error(404)
 
-        with pytest.raises(ServerResponseError) as ctx:
+        with pytest.raises(SystemExit) as ctx:
             dag_command.state(
-                self.parser.parse_args(["dags", "state", self.dag_id, "missing_run"]),
+                self.parser.parse_args(
+                    ["dags", "state", self.dag_id, "--logical-date", "2025-01-01T00:00:00+00:00"]
+                ),
                 api_client=api_client,
             )
 
-        assert ctx.value is error
+        assert ctx.value.code == 1
+        api_client.dag_runs.get.assert_not_called()
+        assert _normalize_rich_output(capsys.readouterr().out) == (
+            "Dag run for test_dag with logical date '2025-01-01T00:00:00+00:00' not found"
+        )
 
-    def test_state_rejects_naive_logical_date(self):
+    def test_state_dag_run_not_found_by_run_id(self, capsys):
         api_client = mock.MagicMock()
-        api_client.dag_runs.get.side_effect = _server_error(404)
+        api_client.dag_runs.get.side_effect = _make_server_error(404)
 
-        with pytest.raises(SystemExit, match="Logical date must include a timezone offset"):
+        with pytest.raises(SystemExit) as ctx:
             dag_command.state(
-                self.parser.parse_args(["dags", "state", self.dag_id, "2025-01-01T00:00:00"]),
+                self.parser.parse_args(["dags", "state", self.dag_id, "missing_run"]),
                 api_client=api_client,
             )
 
+        assert ctx.value.code == 1
         api_client.dag_runs.list.assert_not_called()
+        assert _normalize_rich_output(capsys.readouterr().out) == (
+            "Dag run 'missing_run' of Dag 'test_dag' not found"
+        )
 
-    def test_state_propagates_non_404_api_error(self):
+    def test_state_reraises_non_404_dag_run_get_error(self):
         api_client = mock.MagicMock()
-        api_client.dag_runs.get.side_effect = error = _server_error(500)
+        api_client.dag_runs.get.side_effect = error = _make_server_error(500)
 
         with pytest.raises(ServerResponseError) as ctx:
             dag_command.state(
@@ -332,3 +365,18 @@ def test_state_propagates_non_404_api_error(self):
 
         assert ctx.value is error
         api_client.dag_runs.list.assert_not_called()
+
+    def test_state_reraises_non_404_dag_run_list_error(self):
+        api_client = mock.MagicMock()
+        api_client.dag_runs.list.side_effect = error = _make_server_error(500)
+
+        with pytest.raises(ServerResponseError) as ctx:
+            dag_command.state(
+                self.parser.parse_args(
+                    ["dags", "state", self.dag_id, "--logical-date", "2025-01-01T00:00:00+00:00"]
+                ),
+                api_client=api_client,
+            )
+
+        assert ctx.value is error
+        api_client.dag_runs.get.assert_not_called()