diff --git a/airflow-core/src/airflow/cli/commands/task_command.py b/airflow-core/src/airflow/cli/commands/task_command.py index 550f8f6451980..af1fb4f71a967 100644 --- a/airflow-core/src/airflow/cli/commands/task_command.py +++ b/airflow-core/src/airflow/cli/commands/task_command.py @@ -23,6 +23,7 @@ import json import logging import os +import re import textwrap from contextlib import redirect_stdout from typing import TYPE_CHECKING, Protocol, cast @@ -600,9 +601,13 @@ def task_clear(args) -> None: dags = get_dags(args.bundle_name, args.dag_id, use_regex=args.dag_regex, from_db=True) if args.task_regex: + try: + task_pattern = re.compile(args.task_regex) + except re.error as e: + raise SystemExit(f"Invalid --task-regex {args.task_regex!r}: {e}") for idx, dag in enumerate(dags): dags[idx] = dag.partial_subset( - task_ids=args.task_regex, + task_ids={task_id for task_id in dag.task_ids if task_pattern.search(task_id)}, include_downstream=args.downstream, include_upstream=args.upstream, ) diff --git a/airflow-core/tests/unit/cli/commands/test_task_command.py b/airflow-core/tests/unit/cli/commands/test_task_command.py index a71fb78ca2f5d..3b6b7532a84ea 100644 --- a/airflow-core/tests/unit/cli/commands/test_task_command.py +++ b/airflow-core/tests/unit/cli/commands/test_task_command.py @@ -810,6 +810,34 @@ def test_task_states_for_dag_run(self): "end_date": ti_end.isoformat(), } + @pytest.mark.parametrize( + ("task_regex", "expected_task_ids"), + [ + ("^runme_[02]$", {"runme_0", "runme_2"}), + ("unme_[02]", {"runme_0", "runme_2"}), + ("runme_0|also_run_this", {"runme_0", "also_run_this"}), + ], + ) + @mock.patch("airflow.cli.commands.task_command.SerializedDAG.clear_dags", autospec=True) + def test_task_clear_task_regex_matches_as_regex(self, mock_clear_dags, task_regex, expected_task_ids): + args = self.parser.parse_args( + ["tasks", "clear", "example_bash_operator", "--task-regex", task_regex, "--yes"] + ) + task_command.task_clear(args) + + (dags,) = mock_clear_dags.call_args.args + assert set(dags[0].task_ids) == expected_task_ids + + @mock.patch("airflow.cli.commands.task_command.SerializedDAG.clear_dags", autospec=True) + def test_task_clear_task_regex_invalid_pattern_exits_cleanly(self, mock_clear_dags): + args = self.parser.parse_args( + ["tasks", "clear", "example_bash_operator", "--task-regex", "runme_[", "--yes"] + ) + with pytest.raises(SystemExit, match=r"Invalid --task-regex 'runme_\['"): + task_command.task_clear(args) + + mock_clear_dags.assert_not_called() + def test_task_states_for_dag_run_when_dag_run_not_exists(self): """ task_states_for_dag_run should return an AirflowException when invalid dag id is passed