From 94f9d0854e38b309f1f0d9871c61ec137a2cb8b5 Mon Sep 17 00:00:00 2001 From: Niko Oliveira Date: Tue, 23 Sep 2025 09:36:23 -0700 Subject: [PATCH 1/2] pass args/kwargs to super in Celery executors This allows the team_name to be passed to the super class. This is low hanging fruit to allow the Celery executor to be used for multi team testing. More changes will be needed to allow the Celery executor to use team-based config, but that will be done at a future time. --- .../providers/celery/executors/celery_executor.py | 4 ++-- .../unit/celery/executors/test_celery_executor.py | 11 +++++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py index 95c4c97f12ce6..6091dc5c124d3 100644 --- a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py +++ b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py @@ -290,8 +290,8 @@ class CeleryExecutor(BaseExecutor): # TODO: TaskSDK: move this type change into BaseExecutor queued_tasks: dict[TaskInstanceKey, workloads.All] # type: ignore[assignment] - def __init__(self): - super().__init__() + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) # Celery doesn't support bulk sending the tasks (which can become a bottleneck on bigger clusters) # so we use a multiprocessing pool to speed this up. diff --git a/providers/celery/tests/unit/celery/executors/test_celery_executor.py b/providers/celery/tests/unit/celery/executors/test_celery_executor.py index 9384ef626ce1f..d1820a6d76f2c 100644 --- a/providers/celery/tests/unit/celery/executors/test_celery_executor.py +++ b/providers/celery/tests/unit/celery/executors/test_celery_executor.py @@ -126,6 +126,17 @@ def test_supports_sentry(self): def test_cli_commands_vended(self): assert CeleryExecutor.get_cli_commands() + def test_celery_executor_init_with_args_kwargs(self): + """Test that CeleryExecutor properly passes args and kwargs to BaseExecutor.""" + parallelism = 50 + team_name = "test_team" + + executor = celery_executor.CeleryExecutor(parallelism, team_name=team_name) + + assert executor.parallelism == parallelism + assert executor.team_name == team_name + assert executor.conf.team_name == team_name + @pytest.mark.backend("mysql", "postgres") def test_exception_propagation(self, caplog): caplog.set_level( From fae7bc16343605092466c3a601a75f1f18d1518a Mon Sep 17 00:00:00 2001 From: Niko Oliveira Date: Tue, 23 Sep 2025 14:27:05 -0700 Subject: [PATCH 2/2] Fix compat tests --- .../unit/celery/executors/test_celery_executor.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/providers/celery/tests/unit/celery/executors/test_celery_executor.py b/providers/celery/tests/unit/celery/executors/test_celery_executor.py index d1820a6d76f2c..771b8b232bd92 100644 --- a/providers/celery/tests/unit/celery/executors/test_celery_executor.py +++ b/providers/celery/tests/unit/celery/executors/test_celery_executor.py @@ -131,11 +131,18 @@ def test_celery_executor_init_with_args_kwargs(self): parallelism = 50 team_name = "test_team" - executor = celery_executor.CeleryExecutor(parallelism, team_name=team_name) + if AIRFLOW_V_3_1_PLUS: + # team_name was added in Airflow 3.1 + executor = celery_executor.CeleryExecutor(parallelism=parallelism, team_name=team_name) + else: + executor = celery_executor.CeleryExecutor(parallelism) assert executor.parallelism == parallelism - assert executor.team_name == team_name - assert executor.conf.team_name == team_name + + if AIRFLOW_V_3_1_PLUS: + # team_name was added in Airflow 3.1 + assert executor.team_name == team_name + assert executor.conf.team_name == team_name @pytest.mark.backend("mysql", "postgres") def test_exception_propagation(self, caplog):