From 52dc407821944f6e156b7b8cce60e23c0871f1b4 Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Wed, 22 Jul 2026 21:50:21 +0900 Subject: [PATCH] Split airflowctl dags state argument into run_id and --logical-date Signed-off-by: PoAn Yang --- .../test_airflowctl_commands.py | 2 +- airflow-ctl/src/airflowctl/ctl/cli_config.py | 15 ++- .../airflowctl/ctl/commands/dag_command.py | 82 +++++++----- .../tests/airflow_ctl/api/test_operations.py | 14 ++ .../ctl/commands/test_dag_command.py | 122 ++++++++++++------ 5 files changed, 159 insertions(+), 76 deletions(-) diff --git a/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py b/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py index 87d9851cbf0c0..3c0b4f323fcc3 100644 --- a/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py +++ b/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", diff --git a/airflow-ctl/src/airflowctl/ctl/cli_config.py b/airflow-ctl/src/airflowctl/ctl/cli_config.py index 47270d9c29c10..df60bc3b769dd 100755 --- a/airflow-ctl/src/airflowctl/ctl/cli_config.py +++ b/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( diff --git a/airflow-ctl/src/airflowctl/ctl/commands/dag_command.py b/airflow-ctl/src/airflowctl/ctl/commands/dag_command.py index f0562891d115c..d2acee3d24a7d 100644 --- a/airflow-ctl/src/airflowctl/ctl/commands/dag_command.py +++ b/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)) diff --git a/airflow-ctl/tests/airflow_ctl/api/test_operations.py b/airflow-ctl/tests/airflow_ctl/api/test_operations.py index f988a6f329a99..9c89a5389433b 100644 --- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py +++ b/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) diff --git a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_dag_command.py b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_dag_command.py index f5966f09c0f28..62c402afb6c11 100644 --- a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_dag_command.py +++ b/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()