From cbb9bd0bd632e0ed20099e325203df669d434f3d Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Mon, 4 Sep 2023 23:42:53 +0200 Subject: [PATCH] Add Celery extra in Breeze when using CeleryExecutor and --use-airflow When using `--executor CeleryExecutor` and `--use-airflow` in Airflow 2.7+ you need to also specify `--airflow-extras celery` to make it works. With this change `--airflow-extras celery` is added automatically when CeleryExecutor is used and when `--use-airflow` is specified (unless you manually specify your own extras. This makes it easier to test RC releases with CeleryExecutor. --- .../commands/developer_commands.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/dev/breeze/src/airflow_breeze/commands/developer_commands.py b/dev/breeze/src/airflow_breeze/commands/developer_commands.py index f30c5b904d4ab..6a7e6ac00d778 100644 --- a/dev/breeze/src/airflow_breeze/commands/developer_commands.py +++ b/dev/breeze/src/airflow_breeze/commands/developer_commands.py @@ -721,12 +721,26 @@ def enter_shell(**kwargs) -> RunCommandResult: if shell_params.backend == "sqlite": get_console().print( - f"\n[warn]backend: sqlite is not " + f"\n[warning]backend: sqlite is not " f"compatible with executor: {shell_params.executor}. " f"Changing the executor to SequentialExecutor.\n" ) shell_params.executor = "SequentialExecutor" + if shell_params.executor == "CeleryExecutor" and shell_params.use_airflow_version: + if shell_params.airflow_extras and "celery" not in shell_params.airflow_extras.split(): + get_console().print( + f"\n[warning]CeleryExecutor requires airflow_extras: celery. " + f"Adding celery to extras: '{shell_params.airflow_extras}'.\n" + ) + shell_params.airflow_extras += ",celery" + elif not shell_params.airflow_extras: + get_console().print( + "\n[warning]CeleryExecutor requires airflow_extras: celery. " + "Setting airflow extras to 'celery'.\n" + ) + shell_params.airflow_extras = "celery" + if shell_params.include_mypy_volume: create_mypy_volume_if_needed() shell_params.print_badge_info()