diff --git a/airflow-core/src/airflow/cli/commands/task_command.py b/airflow-core/src/airflow/cli/commands/task_command.py index e2a1ab28e0495..b8cae95ae8087 100644 --- a/airflow-core/src/airflow/cli/commands/task_command.py +++ b/airflow-core/src/airflow/cli/commands/task_command.py @@ -493,6 +493,7 @@ def task_render(args, dag: DAG | None = None) -> None: ) +@deprecated_for_airflowctl("airflowctl tasks clear") @cli_utils.action_cli(check_db=False) @providers_configuration_loaded def task_clear(args) -> None: diff --git a/airflow-core/tests/unit/cli/commands/test_command_deprecations.py b/airflow-core/tests/unit/cli/commands/test_command_deprecations.py index 65982d65adb6f..59d26a12864cb 100644 --- a/airflow-core/tests/unit/cli/commands/test_command_deprecations.py +++ b/airflow-core/tests/unit/cli/commands/test_command_deprecations.py @@ -76,6 +76,7 @@ (config_command.get_value, "airflowctl config get"), (config_command.show_config, "airflowctl config list"), (task_command.task_states_for_dag_run, "airflowctl tasks states-for-dag-run"), + (task_command.task_clear, "airflowctl tasks clear"), ] 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..d2440f05d2466 100644 --- a/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py +++ b/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py @@ -99,6 +99,7 @@ def date_param(): # Tasks commands 'tasks states-for-dag-run example_bash_operator "manual__{date_param}"', 'tasks states-for-dag-run example_bash_operator --logical-date "{date_param}"', + 'tasks clear example_bash_operator --dag-run-id "manual__{date_param}" --task-ids runme_0 -o json', # Task Instances commands 'taskinstances list example_bash_operator "manual__{date_param}"', # XCom commands - need a Dag run with completed tasks diff --git a/airflow-ctl/docs/images/command_hashes.txt b/airflow-ctl/docs/images/command_hashes.txt index 10e59c5959f8b..74b5ede7ce3a2 100644 --- a/airflow-ctl/docs/images/command_hashes.txt +++ b/airflow-ctl/docs/images/command_hashes.txt @@ -1,4 +1,4 @@ -main:0460d9c03248bee26207b20b05aa36b9 +main:2c358f5fc894541cab11854dcd67658a assets:6419e20452692f577c4c6f570b74be0c auth:d79e9c7d00c432bdbcbc2a86e2e32053 backfill:74c8737b0a62a86ed3605fa9e6165874 @@ -10,7 +10,7 @@ jobs:a5b644c5da8889443bb40ee10b599270 pools:19efe105b9515ab1926ebcaf0e028d71 providers:34502fe09dc0b8b0a13e7e46efdffda6 taskinstances:7e323968c0b585287c2a4ab4339aee1f -tasks:089b19625c893d189b5fc02b3abd547e +tasks:e2ed0ab7f67ebacf4be87b5f396ff782 variables:f8fc76d3d398b2780f4e97f7cd816646 version:31f4efdf8de0dbaaa4fac71ff7efecc3 plugins:4864fd8f356704bd2b3cd1aec3567e35 diff --git a/airflow-ctl/docs/images/output_main.svg b/airflow-ctl/docs/images/output_main.svg index 9c8cec5269b68..2a39e7e0ce2dd 100644 --- a/airflow-ctl/docs/images/output_main.svg +++ b/airflow-ctl/docs/images/output_main.svg @@ -19,114 +19,114 @@ font-weight: 700; } - .terminal-3519450412-matrix { + .terminal-151424782-matrix { font-family: Fira Code, monospace; font-size: 20px; line-height: 24.4px; font-variant-east-asian: full-width; } - .terminal-3519450412-title { + .terminal-151424782-title { font-size: 18px; font-weight: bold; font-family: arial; } - .terminal-3519450412-r1 { fill: #ff8700 } -.terminal-3519450412-r2 { fill: #c5c8c6 } -.terminal-3519450412-r3 { fill: #808080 } -.terminal-3519450412-r4 { fill: #68a0b3 } + .terminal-151424782-r1 { fill: #ff8700 } +.terminal-151424782-r2 { fill: #c5c8c6 } +.terminal-151424782-r3 { fill: #808080 } +.terminal-151424782-r4 { fill: #68a0b3 } - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + @@ -138,39 +138,39 @@ - + - - Usage:airflowctl [-h] GROUP_OR_COMMAND... - -Positional Arguments: -GROUP_OR_COMMAND - -    Groups -assetsPerform Assets operations -authManage authentication for CLI. Either pass token from -environment variable/parameter or pass username and -password. -backfillPerform Backfill operations -configPerform Config operations -connectionsPerform Connections operations -dagrunPerform DagRun operations -dagsPerform Dags operations -jobsPerform Jobs operations -pluginsPerform Plugins operations -poolsPerform Pools operations -providersPerform Providers operations -taskinstances -Perform TaskInstances operations -tasksManage Airflow tasks -variablesPerform Variables operations -xcomPerform XCom operations - -    Commands: -versionShow version information - -Options: --h, --helpshow this help message and exit + + Usage:airflowctl [-h] GROUP_OR_COMMAND... + +Positional Arguments: +GROUP_OR_COMMAND + +    Groups +assetsPerform Assets operations +authManage authentication for CLI. Either pass token from +environment variable/parameter or pass username and +password. +backfillPerform Backfill operations +configPerform Config operations +connectionsPerform Connections operations +dagrunPerform DagRun operations +dagsPerform Dags operations +jobsPerform Jobs operations +pluginsPerform Plugins operations +poolsPerform Pools operations +providersPerform Providers operations +taskinstances +Perform TaskInstances operations +tasksPerform Tasks operations +variablesPerform Variables operations +xcomPerform XCom operations + +    Commands: +versionShow version information + +Options: +-h, --helpshow this help message and exit diff --git a/airflow-ctl/docs/images/output_tasks.svg b/airflow-ctl/docs/images/output_tasks.svg index d9349c117529b..f3fa3c4d94506 100644 --- a/airflow-ctl/docs/images/output_tasks.svg +++ b/airflow-ctl/docs/images/output_tasks.svg @@ -1,4 +1,4 @@ - + - - + + - + - + - + - + - + - + - + - + - + - + + + + - + - + - - Usage:airflowctl tasks [-h] COMMAND... - -Manage Airflow tasks - -Positional Arguments: -COMMAND -states-for-dag-run -Get the status of all task instances in a Dag run - -Options: --h, --helpshow this help message and exit + + Usage:airflowctl tasks [-h] COMMAND... + +Perform Tasks operations + +Positional Arguments: +COMMAND +clearClear task instances of a Dag by its ID +states-for-dag-run +Get the status of all task instances in a Dag run + +Options: +-h, --helpshow this help message and exit diff --git a/airflow-ctl/src/airflowctl/api/client.py b/airflow-ctl/src/airflowctl/api/client.py index 7baf87684dd6c..65f90302868e5 100644 --- a/airflow-ctl/src/airflowctl/api/client.py +++ b/airflow-ctl/src/airflowctl/api/client.py @@ -58,6 +58,7 @@ ProvidersOperations, ServerResponseError, TaskInstancesOperations, + TasksOperations, VariablesOperations, VersionOperations, XComOperations, @@ -457,6 +458,12 @@ def task_instances(self): """Operations related to task instances.""" return TaskInstancesOperations(self) + @lru_cache() # type: ignore[prop-decorator] + @property + def tasks(self): + """Operations related to tasks.""" + return TasksOperations(self) + @lru_cache() # type: ignore[prop-decorator] @property def variables(self): diff --git a/airflow-ctl/src/airflowctl/api/operations.py b/airflow-ctl/src/airflowctl/api/operations.py index 44bfa30125fb7..5684dfac8403f 100644 --- a/airflow-ctl/src/airflowctl/api/operations.py +++ b/airflow-ctl/src/airflowctl/api/operations.py @@ -39,6 +39,7 @@ BulkBodyPoolBody, BulkBodyVariableBody, BulkResponse, + ClearTaskInstancesBody, Config, ConnectionBody, ConnectionCollectionResponse, @@ -787,6 +788,23 @@ def list(self, dag_id: str, dag_run_id: str) -> TaskInstanceCollectionResponse | ) +class TasksOperations(BaseOperations): + """Tasks operations.""" + + def clear( + self, dag_id: str, clear_task_instances: ClearTaskInstancesBody + ) -> TaskInstanceCollectionResponse | ServerResponseError: + """Clear task instances of a Dag; with dry_run (the default) only previews the affected task instances.""" + try: + self.response = self.client.post( + f"dags/{dag_id}/clearTaskInstances", + json=clear_task_instances.model_dump(mode="json", exclude_none=True), + ) + return TaskInstanceCollectionResponse.model_validate_json(self.response.content) + except ServerResponseError as e: + raise e + + class VariablesOperations(BaseOperations): """Variable operations.""" diff --git a/airflow-ctl/src/airflowctl/ctl/cli_config.py b/airflow-ctl/src/airflowctl/ctl/cli_config.py index 47270d9c29c10..fbaff969ccaab 100755 --- a/airflow-ctl/src/airflowctl/ctl/cli_config.py +++ b/airflow-ctl/src/airflowctl/ctl/cli_config.py @@ -426,7 +426,21 @@ def __init__(self, file_path: str | Path | None = None): # Exclude parameters that are not needed for CLI from datamodels self.excluded_parameters = ["schema_"] # This list is used to determine if the command/operation needs to output data - self.output_command_list = ["list", "get", "create", "delete", "update", "trigger", "add", "edit"] + self.output_command_list = [ + "list", + "get", + "create", + "delete", + "update", + "trigger", + "add", + "edit", + "clear", + ] + # Datamodels whose generated bool flags follow the datamodel field defaults instead of + # defaulting to False, so the CLI keeps the API semantics (e.g. a bare ``tasks clear`` + # must keep ``dry_run=True`` and preview instead of clearing). + self.field_bool_default_datamodels = ["ClearTaskInstancesBody"] self.exclude_operation_names = ["LoginOperations", "VersionOperations", "BaseOperations"] self.exclude_method_names = [ "error", @@ -593,6 +607,14 @@ def _create_positional_arg( help=arg_help, ) + def _get_bool_arg_default(self, parameter_type: str, field_default: Any) -> bool | None: + """Get default for a generated bool flag: the datamodel field default for datamodels in ``field_bool_default_datamodels``, otherwise False.""" + if parameter_type in self.field_bool_default_datamodels and ( + field_default is None or isinstance(field_default, bool) + ): + return field_default + return False + def _create_arg_for_non_primitive_type( self, parameter_type: str, @@ -608,30 +630,24 @@ def _create_arg_for_non_primitive_type( continue self.datamodels_extended_map[parameter_type].append(field) if type(field_type.annotation) is type: - commands.append( - self._create_arg( - arg_flags=("--" + self._sanitize_arg_parameter_key(field),), - arg_type=self._python_type_from_string(field_type.annotation), - arg_action=argparse.BooleanOptionalAction if field_type.annotation is bool else None, # type: ignore - arg_help=f"{field} for {parameter_key} operation", - arg_default=False if field_type.annotation is bool else None, - ) - ) + annotation = field_type.annotation else: try: annotation = field_type.annotation.__args__[0] except AttributeError: annotation = field_type.annotation - commands.append( - self._create_arg( - arg_flags=("--" + self._sanitize_arg_parameter_key(field),), - arg_type=self._python_type_from_string(annotation), - arg_action=argparse.BooleanOptionalAction if annotation is bool else None, # type: ignore - arg_help=f"{field} for {parameter_key} operation", - arg_default=False if annotation is bool else None, - ) + commands.append( + self._create_arg( + arg_flags=(f"--{self._sanitize_arg_parameter_key(field)}",), + arg_type=self._python_type_from_string(annotation), + arg_action=argparse.BooleanOptionalAction if annotation is bool else None, # type: ignore + arg_help=f"{field} for {parameter_key} operation", + arg_default=self._get_bool_arg_default(parameter_type, field_type.default) + if annotation is bool + else None, ) + ) return commands def _create_args_map_from_operation(self): @@ -711,6 +727,20 @@ def _apply_datamodel_defaults(self, datamodel: type, params: dict) -> dict: ): params["logical_date"] = datetime.datetime.now(datetime.timezone.utc) + # Handle ClearTaskInstancesBody: --task-ids arrives as a single string but the API expects + # a list of task_id or [task_id, map_index]; accept comma-separated ids or a JSON list + if datamodel.__name__ == "ClearTaskInstancesBody" and isinstance(params.get("task_ids"), str): + raw_task_ids = params["task_ids"] + if raw_task_ids.lstrip().startswith("["): + try: + params["task_ids"] = json.loads(raw_task_ids) + except json.JSONDecodeError as e: + raise SystemExit(f"Invalid JSON list for --task-ids {raw_task_ids!r}: {e}") + else: + params["task_ids"] = [ + task_id.strip() for task_id in raw_task_ids.split(",") if task_id.strip() + ] + return params def _create_func_map_from_operation(self): diff --git a/airflow-ctl/src/airflowctl/ctl/help_texts.yaml b/airflow-ctl/src/airflowctl/ctl/help_texts.yaml index 1c9b23d3f9491..9c2c591417b14 100644 --- a/airflow-ctl/src/airflowctl/ctl/help_texts.yaml +++ b/airflow-ctl/src/airflowctl/ctl/help_texts.yaml @@ -88,6 +88,9 @@ providers: taskinstances: list: "List all task instances for a given Dag run" +tasks: + clear: "Clear task instances of a Dag by its ID" + variables: get: "Retrieve a variable by its key" list: "List all variables" diff --git a/airflow-ctl/tests/airflow_ctl/api/test_operations.py b/airflow-ctl/tests/airflow_ctl/api/test_operations.py index f988a6f329a99..aaacdd68696ed 100644 --- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py +++ b/airflow-ctl/tests/airflow_ctl/api/test_operations.py @@ -48,6 +48,7 @@ BulkCreateActionPoolBody, BulkCreateActionVariableBody, BulkResponse, + ClearTaskInstancesBody, Config, ConfigOption, ConfigSection, @@ -1594,6 +1595,49 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert response == self.task_instance_collection_response +class TestTasksOperations: + dag_id = "dag_id" + clear_task_instances = ClearTaskInstancesBody( + dry_run=True, + task_ids=["task_1"], + dag_run_id="dag_run_id", + ) + task_instance_response = TaskInstanceResponse( + id=uuid.UUID("0195c8b0-8bb5-7fc1-9ac9-04c286c47b7f"), + task_id="task_1", + dag_id=dag_id, + dag_run_id="dag_run_id", + map_index=-1, + run_after=datetime.datetime(2025, 1, 1, 0, 0, 0), + try_number=1, + max_tries=2, + task_display_name="task_1", + dag_display_name="dag_display_name", + pool="default_pool", + pool_slots=1, + executor_config="{}", + ) + task_instance_collection_response = TaskInstanceCollectionResponse( + task_instances=[task_instance_response], + total_entries=1, + ) + + def test_clear(self): + expected_body = self.clear_task_instances.model_dump(mode="json", exclude_none=True) + + def handle_request(request: httpx.Request) -> httpx.Response: + assert request.url.path == f"/api/v2/dags/{self.dag_id}/clearTaskInstances" + assert request.headers.get("content-type", "").startswith("application/json") + assert json.loads(request.content.decode()) == expected_body + return httpx.Response( + 200, json=json.loads(self.task_instance_collection_response.model_dump_json()) + ) + + client = make_api_client(transport=httpx.MockTransport(handle_request)) + response = client.tasks.clear(self.dag_id, self.clear_task_instances) + assert response == self.task_instance_collection_response + + class TestVariablesOperations: key = "key" value = "val" diff --git a/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py b/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py index b843b3b75682f..d568cacc0981c 100644 --- a/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py +++ b/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py @@ -25,6 +25,7 @@ import httpx import pytest +from airflowctl.api.datamodels.generated import ClearTaskInstancesBody from airflowctl.api.operations import ServerResponseError from airflowctl.ctl.cli_config import ( ARG_AUTH_TOKEN, @@ -779,12 +780,60 @@ def test_apply_datamodel_defaults_other_datamodel(self): # Should return params unchanged for other datamodels assert result == params, "Params should be unchanged for non-TriggerDAGRunPostBody datamodels" + def test_tasks_clear_args_follow_datamodel_defaults(self): + """Bool flags of ``tasks clear`` keep the ClearTaskInstancesBody defaults so a bare invocation only dry-runs.""" + command_factory = CommandFactory() + tasks_group = next( + group_command for group_command in command_factory.group_commands if group_command.name == "tasks" + ) + clear_command = next( + sub_command for sub_command in tasks_group.subcommands if sub_command.name == "clear" + ) + args_by_flag = {arg.flags[0]: arg for arg in clear_command.args} + + assert "dag_id" in args_by_flag, "required path parameter should be positional" + assert args_by_flag["--dry-run"].kwargs["action"] == BooleanOptionalAction + assert args_by_flag["--dry-run"].kwargs["default"] is True + assert args_by_flag["--only-failed"].kwargs["default"] is True + assert args_by_flag["--reset-dag-runs"].kwargs["default"] is True + assert args_by_flag["--only-running"].kwargs["default"] is False + assert args_by_flag["--run-on-latest-version"].kwargs["default"] is None + assert args_by_flag["--task-ids"].kwargs["type"] is str + assert "--output" in args_by_flag + + @pytest.mark.parametrize( + ("raw_task_ids", "expected_task_ids"), + [ + ("task_1", ["task_1"]), + ("task_1,task_2", ["task_1", "task_2"]), + (" task_1 , task_2 ,", ["task_1", "task_2"]), + ('["task_1", ["mapped_task", 0]]', ["task_1", ["mapped_task", 0]]), + (None, None), + ], + ) + def test_apply_datamodel_defaults_clear_task_instances_task_ids(self, raw_task_ids, expected_task_ids): + """Test _apply_datamodel_defaults parses --task-ids strings for ClearTaskInstancesBody.""" + command_factory = CommandFactory() + result = command_factory._apply_datamodel_defaults( + ClearTaskInstancesBody, {"task_ids": raw_task_ids, "dry_run": True} + ) + + assert result["task_ids"] == expected_task_ids + assert result["dry_run"] is True + + def test_apply_datamodel_defaults_clear_task_instances_invalid_json_task_ids(self): + """Test _apply_datamodel_defaults rejects malformed JSON lists passed to --task-ids.""" + command_factory = CommandFactory() + with pytest.raises(SystemExit, match="Invalid JSON list for --task-ids"): + command_factory._apply_datamodel_defaults(ClearTaskInstancesBody, {"task_ids": '["oops'}) + @pytest.mark.parametrize( ("group_name", "subcommand_name", "expected_help"), [ ("assets", "get", "Retrieve an asset by its ID"), ("connections", "get", "Retrieve a connection by its ID"), ("taskinstances", "list", "List all task instances for a given Dag run"), + ("tasks", "clear", "Clear task instances of a Dag by its ID"), ], ) def test_help_texts_used_for_auto_generated_commands(self, group_name, subcommand_name, expected_help):