This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 76943ec8e44 Fix Celery result backend compatibility with PgBouncer
(#74841)
76943ec8e44 is described below
commit 76943ec8e44f69c7d7af26b0848f8d12f2b7c971
Author: Aaron Chen <[email protected]>
AuthorDate: Sat Oct 17 15:48:11 2026 -0700
Fix Celery result backend compatibility with PgBouncer (#74841)
* Fix Celery result backend compatibility with PgBouncer
* Refactor default Celery executor to improve type casting for result
backend options
---
.../tests/unit/always/test_project_structure.py | 1 -
providers/celery/docs/celery_executor.rst | 17 ++
providers/celery/provider.yaml | 2 +
.../providers/celery/executors/default_celery.py | 30 +++-
.../airflow/providers/celery/get_provider_info.py | 2 +-
.../unit/celery/executors/test_celery_executor.py | 7 +-
.../unit/celery/executors/test_default_celery.py | 175 +++++++++++++++++++++
7 files changed, 228 insertions(+), 6 deletions(-)
diff --git a/airflow-core/tests/unit/always/test_project_structure.py
b/airflow-core/tests/unit/always/test_project_structure.py
index bf80379664c..514e3b16cde 100644
--- a/airflow-core/tests/unit/always/test_project_structure.py
+++ b/airflow-core/tests/unit/always/test_project_structure.py
@@ -72,7 +72,6 @@ class TestProjectStructure:
"providers/amazon/tests/unit/amazon/aws/sensors/test_emr.py",
"providers/amazon/tests/unit/amazon/aws/sensors/test_sagemaker.py",
"providers/celery/tests/unit/celery/executors/test_celery_executor_utils.py",
-
"providers/celery/tests/unit/celery/executors/test_default_celery.py",
"providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor_utils.py",
"providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_exceptions.py",
"providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kube_config.py",
diff --git a/providers/celery/docs/celery_executor.rst
b/providers/celery/docs/celery_executor.rst
index 31a1619ffb0..d175908549d 100644
--- a/providers/celery/docs/celery_executor.rst
+++ b/providers/celery/docs/celery_executor.rst
@@ -88,6 +88,23 @@ Some caveats:
- Tasks can consume resources. Make sure your worker has enough resources to
run ``worker_concurrency`` tasks
- Queue names are limited to 256 characters, but each broker backend might
have its own restrictions
+PostgreSQL result backends
+--------------------------
+
+For result backends using psycopg (v3), Airflow defaults
``connect_args.prepare_threshold`` to
+``None`` to disable automatic prepared statements. This supports a
transaction-mode connection
+pooler that does not track prepared statements. The setting also works when
connecting directly
+to PostgreSQL. Other database drivers, including psycopg2, do not receive this
parameter.
+
+An explicit value in ``[celery] result_backend_sqlalchemy_engine_options``
takes precedence.
+For example, when connecting directly to PostgreSQL or using a pooler
configured to support
+prepared statements, you can restore psycopg's default preparation threshold:
+
+.. code-block:: ini
+
+ [celery]
+ result_backend_sqlalchemy_engine_options = {"connect_args":
{"prepare_threshold": 5}}
+
Redis maintenance
-----------------
diff --git a/providers/celery/provider.yaml b/providers/celery/provider.yaml
index 80277f56fc1..4d93fa78870 100644
--- a/providers/celery/provider.yaml
+++ b/providers/celery/provider.yaml
@@ -275,6 +275,8 @@ config:
result_backend_sqlalchemy_engine_options:
description: |
Optional configuration dictionary to pass to the Celery result
backend SQLAlchemy engine.
+ For psycopg (v3), ``connect_args.prepare_threshold`` defaults to
``None`` to disable
+ automatic prepared statements unless explicitly overridden.
version_added: ~
type: string
example: '{"pool_recycle": 1800}'
diff --git
a/providers/celery/src/airflow/providers/celery/executors/default_celery.py
b/providers/celery/src/airflow/providers/celery/executors/default_celery.py
index 8cc8a0ca445..e57719f070e 100644
--- a/providers/celery/src/airflow/providers/celery/executors/default_celery.py
+++ b/providers/celery/src/airflow/providers/celery/executors/default_celery.py
@@ -23,7 +23,10 @@ import json
import logging
import re
import ssl
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
+
+from sqlalchemy.engine import make_url
+from sqlalchemy.exc import ArgumentError, NoSuchModuleError
from airflow.exceptions import AirflowConfigException as
CoreAirflowConfigException
from airflow.providers.celery.version_compat import AIRFLOW_V_3_0_PLUS
@@ -112,6 +115,23 @@ def _broker_transport_options(broker_url: str, conf:
AirflowSDKConfigParser | An
return broker_transport_options
+def _get_database_engine_options(result_backend: str, options: dict[str, Any])
-> dict[str, Any]:
+ if not result_backend.startswith("db+"):
+ return options
+ try:
+ driver = make_url(result_backend.removeprefix("db+")).get_driver_name()
+ except (ArgumentError, NoSuchModuleError):
+ # Defer invalid or unavailable dialects to Celery's backend
initialization.
+ return options
+ if driver != "psycopg":
+ return options
+ # Celery creates its own engines, so metadata-engine defaults do not
protect transaction poolers.
+ return {
+ **options,
+ "connect_args": {"prepare_threshold": None,
**options.get("connect_args", {})},
+ }
+
+
def get_default_celery_config(team_conf: AirflowSDKConfigParser | Any) ->
dict[str, Any]:
"""
Build Celery configuration using team-aware config.
@@ -177,8 +197,12 @@ def get_default_celery_config(team_conf:
AirflowSDKConfigParser | Any) -> dict[s
),
"result_backend": result_backend,
"result_backend_transport_options": result_backend_transport_options,
- "database_engine_options": team_conf.getjson(
- "celery", "result_backend_sqlalchemy_engine_options", fallback={}
+ "database_engine_options": _get_database_engine_options(
+ result_backend,
+ cast(
+ "dict[str, Any]",
+ team_conf.getjson("celery",
"result_backend_sqlalchemy_engine_options", fallback={}),
+ ),
),
"worker_concurrency": team_conf.getint("celery", "WORKER_CONCURRENCY",
fallback=16),
"worker_enable_remote_control": team_conf.getboolean(
diff --git a/providers/celery/src/airflow/providers/celery/get_provider_info.py
b/providers/celery/src/airflow/providers/celery/get_provider_info.py
index ec9be555b00..f5ecc04ddb9 100644
--- a/providers/celery/src/airflow/providers/celery/get_provider_info.py
+++ b/providers/celery/src/airflow/providers/celery/get_provider_info.py
@@ -141,7 +141,7 @@ def get_provider_info():
"default": None,
},
"result_backend_sqlalchemy_engine_options": {
- "description": "Optional configuration dictionary to
pass to the Celery result backend SQLAlchemy engine.\n",
+ "description": "Optional configuration dictionary to
pass to the Celery result backend SQLAlchemy engine.\nFor psycopg (v3),
``connect_args.prepare_threshold`` defaults to ``None`` to disable\nautomatic
prepared statements unless explicitly overridden.\n",
"version_added": None,
"type": "string",
"example": '{"pool_recycle": 1800}',
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 39d7bf0175b..2247944d4b6 100644
--- a/providers/celery/tests/unit/celery/executors/test_celery_executor.py
+++ b/providers/celery/tests/unit/celery/executors/test_celery_executor.py
@@ -492,7 +492,12 @@ class TestCeleryExecutor:
assert not executor.has_task(ti)
mock_fail.assert_not_called()
- @conf_vars({("celery", "result_backend_sqlalchemy_engine_options"):
'{"pool_recycle": 1800}'})
+ @conf_vars(
+ {
+ ("celery", "result_backend"): "db+sqlite:///:memory:",
+ ("celery", "result_backend_sqlalchemy_engine_options"):
'{"pool_recycle": 1800}',
+ }
+ )
def test_result_backend_sqlalchemy_engine_options(self):
import importlib
diff --git
a/providers/celery/tests/unit/celery/executors/test_default_celery.py
b/providers/celery/tests/unit/celery/executors/test_default_celery.py
new file mode 100644
index 00000000000..77aab202f1c
--- /dev/null
+++ b/providers/celery/tests/unit/celery/executors/test_default_celery.py
@@ -0,0 +1,175 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import json
+
+import pytest
+from celery import current_app
+from celery.contrib.testing.app import setup_default_app
+from sqlalchemy import event
+from sqlalchemy.engine import make_url
+
+from airflow.providers.celery.executors import default_celery
+from airflow.providers.celery.executors.celery_executor_utils import
create_celery_app
+from airflow.providers.common.compat.sdk import conf
+
+from tests_common.test_utils.config import conf_vars
+
+pytestmark = pytest.mark.skipif(
+ not default_celery._USE_PSYCOPG3,
+ reason="Psycopg3 result backends require psycopg and SQLAlchemy 2.0+",
+)
+
+
+class ConnectionAttemptIntercepted(Exception):
+ pass
+
+
[email protected]
+def restore_current_celery_app():
+ with setup_default_app(current_app._get_current_object()):
+ yield
+
+
+def get_backend_connect_kwargs(backend):
+ engine = backend.session_manager.get_engine(backend.url,
**backend.engine_options)
+ connect_kwargs = {}
+
+ @event.listens_for(engine, "do_connect")
+ def capture_connect_kwargs(dialect, connection_record, args, kwargs):
+ connect_kwargs.update(kwargs)
+ raise ConnectionAttemptIntercepted
+
+ try:
+ with pytest.raises(ConnectionAttemptIntercepted):
+ engine.connect()
+ return engine.dialect.driver, connect_kwargs
+ finally:
+ engine.dispose()
+
+
[email protected]("restore_current_celery_app")
[email protected](
+ "result_backend",
+ [
+ pytest.param(None, id="derived-from-metadata"),
+ pytest.param("db+postgresql+psycopg://user:pass@host/db",
id="explicit-psycopg"),
+ pytest.param("db+postgresql://user:pass@host/db",
id="sqlalchemy-default-driver"),
+ pytest.param("db+postgresql+psycopg2://user:pass@host/db",
id="explicit-psycopg2"),
+ ],
+)
[email protected](
+ "engine_options",
+ [
+ pytest.param({}, id="default-options"),
+ pytest.param({"pool_recycle": 1800, "connect_args":
{"connect_timeout": 7}}, id="custom-options"),
+ ],
+)
+def test_result_backend_disables_psycopg_preparation(result_backend,
engine_options):
+ expected_url = result_backend or
"db+postgresql+psycopg://user:pass@host/db"
+ expected_driver =
make_url(expected_url.removeprefix("db+")).get_driver_name()
+ expected_connect_args = {
+ **({"prepare_threshold": None} if expected_driver == "psycopg" else
{}),
+ **engine_options.get("connect_args", {}),
+ }
+ expected_options = {**engine_options}
+ if expected_connect_args:
+ expected_options["connect_args"] = expected_connect_args
+
+ with conf_vars(
+ {
+ # --skip-db-tests sets the deprecated [core] key, which takes
precedence over [database].
+ ("core", "sql_alchemy_conn"): None,
+ ("database", "sql_alchemy_conn"): "postgresql://user:pass@host/db",
+ ("celery", "result_backend"): result_backend,
+ ("celery", "result_backend_sqlalchemy_engine_options"):
json.dumps(engine_options),
+ ("celery", "extra_celery_config"):
'{"database_create_tables_at_setup": false}',
+ }
+ ):
+ with create_celery_app(conf) as app:
+ backend = app.backend
+ driver, connect_kwargs = get_backend_connect_kwargs(backend)
+ assert driver == expected_driver
+ assert backend.engine_options == expected_options
+ assert {key: connect_kwargs[key] for key in expected_connect_args}
== expected_connect_args
+ if expected_driver != "psycopg":
+ assert "prepare_threshold" not in connect_kwargs
+
+
[email protected]("restore_current_celery_app")
[email protected]("prepare_threshold", [None, 0, 5])
+def
test_result_backend_preserves_explicit_prepare_threshold(prepare_threshold):
+ options = {
+ "pool_recycle": 1800,
+ "connect_args": {"prepare_threshold": prepare_threshold,
"connect_timeout": 7},
+ }
+ with conf_vars(
+ {
+ ("celery", "result_backend"):
"db+postgresql+psycopg://user:pass@host/db",
+ ("celery", "result_backend_sqlalchemy_engine_options"):
json.dumps(options),
+ ("celery", "extra_celery_config"):
'{"database_create_tables_at_setup": false}',
+ }
+ ):
+ with create_celery_app(conf) as app:
+ backend = app.backend
+ driver, connect_kwargs = get_backend_connect_kwargs(backend)
+ assert driver == "psycopg"
+ assert backend.engine_options == options
+ assert connect_kwargs["prepare_threshold"] == prepare_threshold
+ assert connect_kwargs["connect_timeout"] == 7
+
+
[email protected]("restore_current_celery_app")
+@conf_vars(
+ {
+ ("celery", "result_backend"): "db+sqlite:///:memory:",
+ ("celery", "result_backend_sqlalchemy_engine_options"): "{}",
+ ("celery", "extra_celery_config"):
'{"database_create_tables_at_setup": false}',
+ }
+)
+def test_result_backend_does_not_add_psycopg_options_to_sqlite():
+ with create_celery_app(conf) as app:
+ backend = app.backend
+ driver, connect_kwargs = get_backend_connect_kwargs(backend)
+ assert driver == "pysqlite"
+ assert backend.engine_options == {}
+ assert "prepare_threshold" not in connect_kwargs
+
+
+@conf_vars(
+ {
+ ("celery", "result_backend"): "redis://redis:6379/0",
+ ("celery", "result_backend_sqlalchemy_engine_options"):
'{"pool_recycle": 1800}',
+ }
+)
+def test_non_database_result_backend_keeps_engine_options():
+ config = default_celery.get_default_celery_config(conf)
+ assert config["result_backend"] == "redis://redis:6379/0"
+ assert config["database_engine_options"] == {"pool_recycle": 1800}
+
+
+@conf_vars(
+ {
+ ("celery", "result_backend"): "db+bad_schema:///",
+ ("celery", "result_backend_sqlalchemy_engine_options"):
'{"pool_recycle": 1800}',
+ }
+)
+def test_unknown_database_dialect_keeps_engine_options():
+ config = default_celery.get_default_celery_config(conf)
+ assert config["result_backend"] == "db+bad_schema:///"
+ assert config["database_engine_options"] == {"pool_recycle": 1800}
