Added unit test for the additional config and addressed comments

apache · Dec 18, 2024 · 84f1e81 · 84f1e81
1 parent 6ad3bc4
commit 84f1e81
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 2 deletions.
diff --git a/chart/values.yaml b/chart/values.yaml
@@ -2607,7 +2607,6 @@ config:
   celery:
     flower_url_prefix: '{{ ternary "" .Values.ingress.flower.path (eq .Values.ingress.flower.path "/") }}'
     worker_concurrency: 16
-    extra_celery_config: '{}'
   scheduler:
     standalone_dag_processor: '{{ ternary "True" "False" .Values.dagProcessor.enabled }}'
     # statsd params included for Airflow 1.10 backward compatibility; moved to [metrics] in 2.0

diff --git a/providers/src/airflow/providers/celery/executors/default_celery.py b/providers/src/airflow/providers/celery/executors/default_celery.py
@@ -69,7 +69,7 @@ def _broker_supports_visibility_timeout(url):
     log.debug("Value for celery result_backend not found. Using sql_alchemy_conn with db+ prefix.")
     result_backend = f'db+{conf.get("database", "SQL_ALCHEMY_CONN")}'
 
-extra_celery_config: dict = conf.getjson("celery", "extra_celery_config", fallback={}) if conf.has_option("celery", "extra_celery_config") else {}
+extra_celery_config: dict = conf.getjson("celery", "extra_celery_config", fallback={})
 
 DEFAULT_CELERY_CONFIG = {
     "accept_content": ["json"],

diff --git a/providers/src/airflow/providers/celery/provider.yaml b/providers/src/airflow/providers/celery/provider.yaml
@@ -329,6 +329,13 @@ config:
         type: string
         example: ~
         default: "False"
+      extra_celery_config:
+        description: |
+          Extra celery configs to include in the celery worker
+        version_added: ~
+        type: string
+        example: ~
+        default: "{}"
   celery_broker_transport_options:
     description: |
       This section is for specifying options which can be passed to the

diff --git a/providers/tests/celery/executors/test_celery_executor.py b/providers/tests/celery/executors/test_celery_executor.py
@@ -399,3 +399,14 @@ def test_celery_task_acks_late_loaded_from_string():
     # reload celery conf to apply the new config
     importlib.reload(default_celery)
     assert default_celery.DEFAULT_CELERY_CONFIG["task_acks_late"] is False
+
+
+@conf_vars({("celery", "extra_celery_config"): '{"worker_max_tasks_per_child": 10}'})
+def test_celery_extra_celery_config_loaded_from_string():
+    import importlib
+
+    # reload celery conf to apply the new config
+    importlib.reload(default_celery)
+    assert default_celery.DEFAULT_CELERY_CONFIG["extra_celery_config"] == {
+        "worker_max_tasks_per_child": 10
+    }