diff --git a/providers/celery/docs/changelog.rst b/providers/celery/docs/changelog.rst index b0650015fa0ab..5d1e48876e85f 100644 --- a/providers/celery/docs/changelog.rst +++ b/providers/celery/docs/changelog.rst @@ -27,6 +27,13 @@ Changelog --------- +.. note:: + When ``[celery] result_backend`` is not set and Airflow derives it from ``sql_alchemy_conn``, a + driverless ``postgresql://`` connection string now produces ``db+postgresql+psycopg://...`` instead + of ``db+postgresql+psycopg2://...``, matching the sync Postgres driver default change in + ``apache-airflow`` and ``apache-airflow-providers-postgres``. An explicit ``[celery] result_backend`` + or an explicit driver in ``sql_alchemy_conn`` is unaffected. + 3.21.0 ...... diff --git a/providers/celery/provider.yaml b/providers/celery/provider.yaml index e10f5030bd470..a5e996df4b597 100644 --- a/providers/celery/provider.yaml +++ b/providers/celery/provider.yaml @@ -247,7 +247,7 @@ config: version_added: ~ type: string sensitive: true - example: "db+postgresql+psycopg2://postgres:airflow@postgres/airflow" + example: "db+postgresql+psycopg://postgres:airflow@postgres/airflow" default: ~ result_backend_sqlalchemy_engine_options: description: | 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 9dd94036734e0..3d90b827d5123 100644 --- a/providers/celery/src/airflow/providers/celery/executors/default_celery.py +++ b/providers/celery/src/airflow/providers/celery/executors/default_celery.py @@ -31,6 +31,23 @@ log = logging.getLogger(__name__) +_USE_PSYCOPG3: bool +try: + from importlib.util import find_spec + + import sqlalchemy + from packaging.version import Version + + _sqlalchemy_version = Version(sqlalchemy.__version__) + _is_sqla2 = (_sqlalchemy_version.major, _sqlalchemy_version.minor, _sqlalchemy_version.micro) >= ( + 2, + 0, + 0, + ) + _USE_PSYCOPG3 = find_spec("psycopg") is not None and _is_sqla2 +except (ImportError, ModuleNotFoundError): + _USE_PSYCOPG3 = False + def _broker_supports_visibility_timeout(url): return url.startswith(("redis://", "rediss://", "sqs://", "sentinel://")) @@ -79,9 +96,12 @@ def get_default_celery_config(team_conf) -> dict[str, Any]: else: log.debug("Value for celery result_backend not found. Using sql_alchemy_conn with db+ prefix.") sql_alchemy_conn = team_conf.get("database", "SQL_ALCHEMY_CONN") - # In SQLAlchemy 2.1 the default PostgreSQL driver changed from psycopg2 to psycopg (v3). - # To maintain existing behavior, we explicitly specify psycopg2 for driverless PostgreSQL URLs. - sql_alchemy_conn = sql_alchemy_conn.replace("postgresql://", "postgresql+psycopg2://", 1) + # Airflow's sync metadata-database driver default is psycopg (v3); mirror that explicitly + # for driverless PostgreSQL URLs instead of relying on SQLAlchemy's own default. Fall back to + # psycopg2 when psycopg/SQLAlchemy 2.0 aren't both available, matching PostgresHook's own + # USE_PSYCOPG3 gating so this doesn't break environments still on the older combination. + target_scheme = "postgresql+psycopg://" if _USE_PSYCOPG3 else "postgresql+psycopg2://" + sql_alchemy_conn = sql_alchemy_conn.replace("postgresql://", target_scheme, 1) result_backend = f"db+{sql_alchemy_conn}" # Handle result backend transport options (for Redis Sentinel support) 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 a715038d84143..4a79cc58821d5 100644 --- a/providers/celery/src/airflow/providers/celery/get_provider_info.py +++ b/providers/celery/src/airflow/providers/celery/get_provider_info.py @@ -130,7 +130,7 @@ def get_provider_info(): "version_added": None, "type": "string", "sensitive": True, - "example": "db+postgresql+psycopg2://postgres:airflow@postgres/airflow", + "example": "db+postgresql+psycopg://postgres:airflow@postgres/airflow", "default": None, }, "result_backend_sqlalchemy_engine_options": { 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 c349efa32cb4e..a4d11a0de0a82 100644 --- a/providers/celery/tests/unit/celery/executors/test_celery_executor.py +++ b/providers/celery/tests/unit/celery/executors/test_celery_executor.py @@ -789,6 +789,32 @@ def test_result_backend_transport_options_with_multiple_options(): assert result_backend_opts["master_name"] == "mymaster" +@conf_vars( + { + ("celery", "result_backend"): None, + ("database", "sql_alchemy_conn"): "postgresql://user:pass@host/db", + } +) +def test_result_backend_derived_from_sql_alchemy_conn_uses_psycopg(monkeypatch): + """A driverless sql_alchemy_conn must derive a psycopg (v3) result_backend, not psycopg2.""" + monkeypatch.setattr(default_celery, "_USE_PSYCOPG3", True) + config = default_celery.get_default_celery_config(conf) + assert config["result_backend"] == "db+postgresql+psycopg://user:pass@host/db" + + +@conf_vars( + { + ("celery", "result_backend"): None, + ("database", "sql_alchemy_conn"): "postgresql://user:pass@host/db", + } +) +def test_result_backend_falls_back_to_psycopg2_without_psycopg3(monkeypatch): + """Without psycopg/SQLAlchemy 2.0 available, the derivation must fall back to psycopg2.""" + monkeypatch.setattr(default_celery, "_USE_PSYCOPG3", False) + config = default_celery.get_default_celery_config(conf) + assert config["result_backend"] == "db+postgresql+psycopg2://user:pass@host/db" + + @conf_vars({("celery_result_backend_transport_options", "sentinel_kwargs"): "invalid_json"}) def test_result_backend_sentinel_kwargs_invalid_json(): """Test that invalid JSON in sentinel_kwargs raises an error."""