Skip to content

Commit bcc430f

Browse files
Eason09053360henry3260
Eason09053360
and
henry3260
authored
Fix DAG.cli() crashing on dags pause and unpause (#72109)
* Fix DAG.cli() crashing on dags pause and unpause A Dag file run as a script goes through DAG.cli(), whose parser drops --dag-id and instead dispatches every subcommand with the Dag object as a second positional argument. The pause and unpause handlers never accepted it, so Dag authors got a TypeError instead of the command. The handlers are only reachable this way from a Dag file, which is why the regular `airflow dags pause` path has always worked and no test covered the difference. * Ignore --treat-dag-id-as-regex when DAG.cli() supplies the Dag The Dag-scoped parser drops --dag-id but keeps --treat-dag-id-as-regex, so the flag survives into a context where the user has no pattern left to supply. The Dag's own id was then read back as a pattern, and dag_ids may contain dots, so an unanchored match could pause unrelated Dags. Normalising the flag next to the dag_id it guards covers both of its readers, including the confirmation prompt, rather than guarding each reader in turn and leaving the next one to be found later. * Update airflow-core/tests/unit/cli/commands/test_dag_command.py --------- Co-authored-by: Eason09053360 <[email protected]> Co-authored-by: Henry Chen <[email protected]>
1 parent 6d4808b commit bcc430f

2 files changed

Lines changed: 38 additions & 6 deletions

File tree

‎airflow-core/src/airflow/cli/commands/dag_command.py‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -243,23 +243,27 @@ def _bulk_clear_runs(
243243
@cli_utils.action_cli
244244
@deprecated_for_airflowctl("airflowctl dags pause")
245245
@providers_configuration_loaded
246-
def dag_pause(args) -> None:
246+
def dag_pause(args, dag: DAG | None = None) -> None:
247247
"""Pauses a DAG."""
248-
set_is_paused(True, args)
248+
set_is_paused(True, args, dag)
249249

250250

251251
@cli_utils.action_cli
252252
@deprecated_for_airflowctl("airflowctl dags unpause")
253253
@providers_configuration_loaded
254-
def dag_unpause(args) -> None:
254+
def dag_unpause(args, dag: DAG | None = None) -> None:
255255
"""Unpauses a DAG."""
256-
set_is_paused(False, args)
256+
set_is_paused(False, args, dag)
257257

258258

259259
@providers_configuration_loaded
260260
@provide_session
261-
def set_is_paused(is_paused: bool, args, *, session: Session = NEW_SESSION) -> None:
261+
def set_is_paused(is_paused: bool, args, dag: DAG | None = None, *, session: Session = NEW_SESSION) -> None:
262262
"""Set is_paused for DAG by a given dag_id."""
263+
if dag:
264+
# A Dag object fully determines the target, so pattern matching has nothing left to match on.
265+
args.dag_id = dag.dag_id
266+
args.treat_dag_id_as_regex = False
263267
query = select(DagModel)
264268
if args.treat_dag_id_as_regex:
265269
query = query.where(DagModel.dag_id.regexp_match(args.dag_id))
@@ -274,7 +278,7 @@ def set_is_paused(is_paused: bool, args, *, session: Session = NEW_SESSION) -> N
274278
return
275279

276280
if not args.yes and args.treat_dag_id_as_regex:
277-
dags_ids = [dag.dag_id for dag in matched_dags]
281+
dags_ids = [dag_model.dag_id for dag_model in matched_dags]
278282
question = (
279283
f"You are about to {'un' if not is_paused else ''}pause {len(dags_ids)} DAGs:\n"
280284
f"{','.join(dags_ids)}"

‎airflow-core/tests/unit/cli/commands/test_dag_command.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -564,6 +564,34 @@ def test_pause(self):
564564
dag_command.dag_unpause(args)
565565
assert not DagModel.get_dagmodel("example_bash_operator").is_paused
566566

567+
def test_pause_unpause_from_dag_cli(self):
568+
"""``DAG.cli()`` passes the Dag positionally and its parser drops ``--dag-id``."""
569+
parser = cli_parser.get_parser(dag_parser=True)
570+
dag = DAG("example_bash_operator")
571+
572+
dag_command.dag_pause(parser.parse_args(["dags", "pause"]), dag)
573+
assert DagModel.get_dagmodel("example_bash_operator").is_paused
574+
575+
dag_command.dag_unpause(parser.parse_args(["dags", "unpause"]), dag)
576+
assert not DagModel.get_dagmodel("example_bash_operator").is_paused
577+
578+
@mock.patch("airflow.cli.commands.dag_command.ask_yesno")
579+
def test_pause_from_dag_cli_ignores_treat_dag_id_as_regex(self, mock_yesno):
580+
"""The Dag fixes the target, so its dag_id must not be read back as a pattern."""
581+
target = DAG("dag.cli_regex_target")
582+
sync_dag_to_db(target)
583+
sync_dag_to_db(DAG("dagXcli_regex_target"))
584+
parser = cli_parser.get_parser(dag_parser=True)
585+
586+
dag_command.dag_pause(parser.parse_args(["dags", "pause", "--treat-dag-id-as-regex"]), target)
587+
588+
mock_yesno.assert_not_called()
589+
assert DagModel.get_dagmodel("dag.cli_regex_target").is_paused
590+
assert not DagModel.get_dagmodel("dagXcli_regex_target").is_paused
591+
592+
clear_db_dags()
593+
self.setup_class()
594+
567595
@mock.patch("airflow.cli.commands.dag_command.ask_yesno")
568596
def test_pause_regex(self, mock_yesno):
569597
args = self.parser.parse_args(["dags", "pause", "^example_.*$", "--treat-dag-id-as-regex"])

0 commit comments

Comments
 (0)