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 924cbde90d598..67c241b440f3c 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/api/operations.py b/airflow-ctl/src/airflowctl/api/operations.py index f58ad78adeb37..7dc22f04c3fdd 100644 --- a/airflow-ctl/src/airflowctl/api/operations.py +++ b/airflow-ctl/src/airflowctl/api/operations.py @@ -634,6 +634,8 @@ def list( logical_date_gte: datetime.datetime | None = None, logical_date_lte: datetime.datetime | None = None, order_by: str | None = None, + *, + suppress_error_log: bool = False, ) -> DAGRunCollectionResponse | ServerResponseError: """ List dag runs (at most `limit` results). @@ -647,6 +649,7 @@ def list( logical_date_gte: Filter dag runs with a logical date greater than or equal to this value. logical_date_lte: Filter dag runs with a logical date less than or equal to this value. order_by: Order the results by the specified field. + suppress_error_log: Skip client-side error logging, for callers handling the error themselves. """ # Use "~" for all DAGs if dag_id is not specified if not dag_id: @@ -663,7 +666,11 @@ def list( ) try: - self.response = self.client.get(f"/dags/{dag_id}/dagRuns", params=params) + self.response = self.client.get( + f"/dags/{dag_id}/dagRuns", + params=params, + extensions={"airflowctl_suppress_error_log": suppress_error_log}, + ) return DAGRunCollectionResponse.model_validate_json(self.response.content) except ServerResponseError as e: raise e diff --git a/airflow-ctl/src/airflowctl/ctl/cli_config.py b/airflow-ctl/src/airflowctl/ctl/cli_config.py index 40ee7a1c67a21..5df8d2744d04d 100755 --- a/airflow-ctl/src/airflowctl/ctl/cli_config.py +++ b/airflow-ctl/src/airflowctl/ctl/cli_config.py @@ -279,12 +279,18 @@ def _load_help_texts_yaml() -> dict[str, dict[str, str]]: ARG_DAG_ID = Arg( flags=("dag_id",), type=str, - help="The Dag ID of the Dag to pause or unpause", + 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)", ) ARG_ACTION_ON_EXISTING_KEY = Arg( @@ -1002,7 +1008,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..b00f8fab3164e 100644 --- a/airflow-ctl/src/airflowctl/ctl/commands/dag_command.py +++ b/airflow-ctl/src/airflowctl/ctl/commands/dag_command.py @@ -108,53 +108,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, 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, 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 dc5cdd396db8b..345fc33e57da5 100644 --- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py +++ b/airflow-ctl/tests/airflow_ctl/api/test_operations.py @@ -1230,6 +1230,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..abd927165a880 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,91 @@ 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, match="1"): + dag_command.state( + self.parser.parse_args(["dags", "state", self.dag_id, *extra_args]), + api_client=api_client, + ) + + 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, match="1"): + dag_command.state( + self.parser.parse_args(["dags", "state", self.dag_id, "--logical-date", logical_date]), + api_client=api_client, + ) + + 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, match="1"): 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 + 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, match="1"): 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, ) 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 +361,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()