kaxil commented on code in PR #74581:
URL: https://github.com/apache/airflow/pull/74581#discussion_r4249431338
##########
airflow-core/docs/howto/set-up-database.rst:
##########
@@ -235,8 +235,10 @@ When ``[database] sql_alchemy_conn_async`` is not set, its
URL is derived from `
postgresql+psycopg_async://<user>:<password>@<host>/<db>
-psycopg3 is the default because it is safe behind transaction-mode PgBouncer
(recommended for all production Postgres installations, see the note below)
with no extra configuration.
-psycopg3 `defers statement preparation
<https://www.psycopg.org/psycopg3/docs/advanced/prepare.html>`_ behind a
threshold (default 5 executions) and lets you disable it entirely with
``prepare_threshold=None``.
+psycopg3 is the default because Airflow can make it safe behind
transaction-mode PgBouncer (recommended for all production Postgres
installations, see the note below) automatically, which it cannot do for
asyncpg without you opting in.
+psycopg3 `prepares statements server-side
<https://www.psycopg.org/psycopg3/docs/advanced/prepare.html>`_ once a
statement has been executed ``prepare_threshold`` times (5 by default), which a
transaction-pooling pooler cannot honor unless it tracks prepared statements
itself.
+Airflow therefore sets ``prepare_threshold = None`` by default for both the
sync and async psycopg3 engines, which disables the behaviour.
+If your pooler tracks prepared statements — PgBouncer has supported them since
1.21 and enables them by default from 1.24, where ``max_prepared_statements``
defaults to 200 — or you run without a pooler at all, you can opt back in by
pointing ``sql_alchemy_connect_args`` (or ``sql_alchemy_connect_args_async``)
at a dictionary that sets ``prepare_threshold``.
Review Comment:
psycopg's [PgBouncer
notes](https://www.psycopg.org/psycopg3/docs/advanced/prepare.html#using-prepared-statements-with-pgbouncer)
put the floor at PgBouncer 1.22, not 1.21, and also need libpq 17 or newer on
the client (or `prepared_max=None`). Someone on 1.21 who opts back in from this
paragraph ends up with the same failure this PR fixes. The same 1.21 is in the
newsfragment (line 7) and the `_psycopg3_connect_args` docstring.
##########
airflow-core/newsfragments/74581.significant.rst:
##########
@@ -0,0 +1,20 @@
+psycopg3 statement preparation is now disabled by default
+
+Airflow 3.4 makes psycopg3 the default Postgres driver for the metadata
database. psycopg3 prepares a statement server-side once it has been executed
``prepare_threshold`` times (5 by default). A pooler running in transaction
mode routes consecutive transactions on one client connection to different
server connections, so unless it tracks prepared statements itself, queries
begin failing with ``InvalidSqlStatementName`` -- or
``DuplicatePreparedStatement`` where several clients generate the same
statement name -- once the threshold is crossed.
+
+Airflow now sets ``prepare_threshold = None`` on both the sync and async
psycopg3 engines, which disables the behaviour. This turns preparation off
entirely rather than only raising the threshold, so Postgres re-parses each
statement and bulk ``executemany`` writes no longer reuse a single prepared
statement.
+
+**Who this affected:** PgBouncer has supported prepared statements since 1.21,
but only enables them by default from 1.24, where ``max_prepared_statements``
defaults to 200. The Airflow Helm chart sets ``pool_mode = transaction`` and
ships PgBouncer 1.23.1 without setting ``max_prepared_statements``, so the
recommended deployment was an affected one. Other poolers may not support
prepared statements at all.
Review Comment:
Neither #68496 nor #69469 is in a release yet (3.3.2 still derives
`asyncpg`), so no released deployment hit this. The people this note changes
things for are those who set `postgresql+psycopg://` or
`postgresql+psycopg_async://` explicitly on 3.3: they silently lose statement
preparation on upgrade. Could the note lead with them rather than the chart
scenario?
##########
airflow-core/tests/unit/core/test_settings.py:
##########
@@ -708,3 +710,68 @@ async def run():
"was never awaited",
):
assert diagnostic not in output, output
+
+
[email protected](
+ ("sql_alchemy_conn", "expected"),
+ [
+ ("postgresql+psycopg://user:pass@host/db", {"prepare_threshold":
None}),
+ ("postgresql+psycopg_async://user:pass@host/db", {"prepare_threshold":
None}),
+ ("postgresql+psycopg:///db?host=/var/run/postgresql",
{"prepare_threshold": None}),
+ ("postgresql+psycopg2://user:pass@host/db", {}),
+ ("postgresql+asyncpg://user:pass@host/db", {}),
+ ("sqlite:////root/airflow.db", {}),
+ ("mysql+aiomysql://user:pass@host/db", {}),
+ # SQL_ALCHEMY_CONN_ASYNC is Optional and unset until configure_orm
runs.
+ (None, {}),
+ ],
+)
+def test_psycopg3_connect_args_only_targets_psycopg3(sql_alchemy_conn,
expected):
+ assert settings._psycopg3_connect_args(sql_alchemy_conn) == expected
+
+
[email protected]("mode", ["sync", "async"])
+def test_get_connect_args_defaults_off_auto_prepare_for_psycopg3(mode):
+ with (
+ patch.object(settings, "SQL_ALCHEMY_CONN",
"postgresql+psycopg://user:pass@host/db"),
+ patch.object(settings, "SQL_ALCHEMY_CONN_ASYNC",
"postgresql+psycopg_async://user:pass@host/db"),
+ ):
+ assert settings._get_connect_args(mode) == {"prepare_threshold": None}
+
+
[email protected](
+ ("mode", "option"),
+ [("sync", "sql_alchemy_connect_args"), ("async",
"sql_alchemy_connect_args_async")],
+)
[email protected](
+ ("configured", "expected"),
+ [
+ ({"prepare_threshold": 5}, {"prepare_threshold": 5}),
+ ({"application_name": "airflow"}, {"prepare_threshold": None,
"application_name": "airflow"}),
+ ],
+ ids=["operator_restores_auto_prepare", "operator_args_merge_over_default"],
+)
+def test_get_connect_args_lets_operator_override_the_psycopg3_default(mode,
option, configured, expected):
+ # conf_vars re-runs configure_vars(), which rebinds SQL_ALCHEMY_CONN -- so
it has to
+ # enter before the patches that set the URL this test depends on.
+ with (
+ conf_vars({("database", option):
"unit.core.test_settings.CONFIGURED_CONNECT_ARGS"}),
Review Comment:
`unit.core.test_settings.CONFIGURED_CONNECT_ARGS` doesn't exist; this passes
only because `getimport` is patched out (without `autospec`).
`test_sqlalchemy_config.py` points the option at a real module-level dict,
which would let you drop the `getimport` patch and exercise the actual import.
##########
airflow-core/tests/unit/core/test_settings.py:
##########
@@ -708,3 +710,68 @@ async def run():
"was never awaited",
):
assert diagnostic not in output, output
+
+
[email protected](
+ ("sql_alchemy_conn", "expected"),
+ [
+ ("postgresql+psycopg://user:pass@host/db", {"prepare_threshold":
None}),
+ ("postgresql+psycopg_async://user:pass@host/db", {"prepare_threshold":
None}),
+ ("postgresql+psycopg:///db?host=/var/run/postgresql",
{"prepare_threshold": None}),
+ ("postgresql+psycopg2://user:pass@host/db", {}),
+ ("postgresql+asyncpg://user:pass@host/db", {}),
+ ("sqlite:////root/airflow.db", {}),
+ ("mysql+aiomysql://user:pass@host/db", {}),
+ # SQL_ALCHEMY_CONN_ASYNC is Optional and unset until configure_orm
runs.
+ (None, {}),
+ ],
+)
+def test_psycopg3_connect_args_only_targets_psycopg3(sql_alchemy_conn,
expected):
+ assert settings._psycopg3_connect_args(sql_alchemy_conn) == expected
+
+
[email protected]("mode", ["sync", "async"])
+def test_get_connect_args_defaults_off_auto_prepare_for_psycopg3(mode):
+ with (
+ patch.object(settings, "SQL_ALCHEMY_CONN",
"postgresql+psycopg://user:pass@host/db"),
+ patch.object(settings, "SQL_ALCHEMY_CONN_ASYNC",
"postgresql+psycopg_async://user:pass@host/db"),
+ ):
+ assert settings._get_connect_args(mode) == {"prepare_threshold": None}
+
+
[email protected](
+ ("mode", "option"),
+ [("sync", "sql_alchemy_connect_args"), ("async",
"sql_alchemy_connect_args_async")],
+)
[email protected](
+ ("configured", "expected"),
+ [
+ ({"prepare_threshold": 5}, {"prepare_threshold": 5}),
+ ({"application_name": "airflow"}, {"prepare_threshold": None,
"application_name": "airflow"}),
+ ],
+ ids=["operator_restores_auto_prepare", "operator_args_merge_over_default"],
+)
+def test_get_connect_args_lets_operator_override_the_psycopg3_default(mode,
option, configured, expected):
+ # conf_vars re-runs configure_vars(), which rebinds SQL_ALCHEMY_CONN -- so
it has to
+ # enter before the patches that set the URL this test depends on.
+ with (
+ conf_vars({("database", option):
"unit.core.test_settings.CONFIGURED_CONNECT_ARGS"}),
+ patch.object(settings, "SQL_ALCHEMY_CONN",
"postgresql+psycopg://user:pass@host/db"),
+ patch.object(settings, "SQL_ALCHEMY_CONN_ASYNC",
"postgresql+psycopg_async://user:pass@host/db"),
+ patch.object(settings.conf, "getimport", return_value=configured),
+ ):
+ assert settings._get_connect_args(mode) == expected
+
+
+def test_psycopg3_default_reaches_the_created_engine():
Review Comment:
This wires `_get_connect_args("sync")` into `create_metadata_engine` by
hand, so it still passes if `configure_orm` or `_configure_async_session` stop
passing the result through. The existing wiring tests only check that a
`connect_args` key is present. Could this go through `configure_orm()` and
`_configure_async_session()` and assert `connect_args == {"prepare_threshold":
None}`? The docs promise the default for both engines, and nothing covers the
async one.
##########
airflow-core/newsfragments/74581.significant.rst:
##########
@@ -0,0 +1,20 @@
+psycopg3 statement preparation is now disabled by default
+
+Airflow 3.4 makes psycopg3 the default Postgres driver for the metadata
database. psycopg3 prepares a statement server-side once it has been executed
``prepare_threshold`` times (5 by default). A pooler running in transaction
mode routes consecutive transactions on one client connection to different
server connections, so unless it tracks prepared statements itself, queries
begin failing with ``InvalidSqlStatementName`` -- or
``DuplicatePreparedStatement`` where several clients generate the same
statement name -- once the threshold is crossed.
+
+Airflow now sets ``prepare_threshold = None`` on both the sync and async
psycopg3 engines, which disables the behaviour. This turns preparation off
entirely rather than only raising the threshold, so Postgres re-parses each
statement and bulk ``executemany`` writes no longer reuse a single prepared
statement.
+
+**Who this affected:** PgBouncer has supported prepared statements since 1.21,
but only enables them by default from 1.24, where ``max_prepared_statements``
defaults to 200. The Airflow Helm chart sets ``pool_mode = transaction`` and
ships PgBouncer 1.23.1 without setting ``max_prepared_statements``, so the
recommended deployment was an affected one. Other poolers may not support
prepared statements at all.
+
+**To keep statement preparation**, if your pooler tracks prepared statements
or you run without one, point ``sql_alchemy_connect_args`` at a dictionary
defined in your ``airflow_local_settings.py``:
+
+.. code-block:: ini
+
+ [database]
+ sql_alchemy_connect_args = airflow_local_settings.PSYCOPG3_CONNECT_ARGS
+
+.. code-block:: python
+
+ PSYCOPG3_CONNECT_ARGS = {"prepare_threshold": 5}
+
+Anything set there takes precedence over the default. The same applies to
``sql_alchemy_connect_args_async`` for the async engine. Operators running
PgBouncer 1.21 or newer can instead set ``max_prepared_statements`` on the
pooler and keep preparation enabled on both sides.
Review Comment:
This reads as if raising `max_prepared_statements` on the pooler is enough
on its own, but `_get_connect_args` still forces `prepare_threshold=None`
whatever the pooler supports, so preparation stays off on the client side. It
has to be both: `max_prepared_statements` on PgBouncer and `prepare_threshold`
restored through `sql_alchemy_connect_args` as shown above.
##########
airflow-core/tests/unit/core/test_settings.py:
##########
@@ -708,3 +710,68 @@ async def run():
"was never awaited",
):
assert diagnostic not in output, output
+
+
[email protected](
+ ("sql_alchemy_conn", "expected"),
+ [
+ ("postgresql+psycopg://user:pass@host/db", {"prepare_threshold":
None}),
+ ("postgresql+psycopg_async://user:pass@host/db", {"prepare_threshold":
None}),
+ ("postgresql+psycopg:///db?host=/var/run/postgresql",
{"prepare_threshold": None}),
+ ("postgresql+psycopg2://user:pass@host/db", {}),
+ ("postgresql+asyncpg://user:pass@host/db", {}),
+ ("sqlite:////root/airflow.db", {}),
+ ("mysql+aiomysql://user:pass@host/db", {}),
+ # SQL_ALCHEMY_CONN_ASYNC is Optional and unset until configure_orm
runs.
Review Comment:
Nit: `SQL_ALCHEMY_CONN_ASYNC` is assigned in `configure_vars`, not
`configure_orm`.
##########
airflow-core/tests/unit/core/test_settings.py:
##########
@@ -35,6 +35,8 @@
from airflow import settings
from airflow.exceptions import AirflowClusterPolicyViolation,
AirflowConfigException
+from tests_common.test_utils.config import conf_vars
Review Comment:
Nit: with this at module level, the local `conf_vars` import inside
`test_explicit_sql_alchemy_conn_async_is_not_rewritten` is now redundant.
##########
airflow-core/src/airflow/settings.py:
##########
@@ -374,14 +374,49 @@ def _is_sqlite_db_path_relative(sqla_conn_str: str) ->
bool:
return True
+def _psycopg3_connect_args(sql_alchemy_conn: str | None) -> dict[str, Any]:
+ """
+ Disable psycopg3's automatic server-side statement preparation.
+
+ psycopg3 prepares a statement server-side after ``prepare_threshold``
executions
+ (5 by default). Transaction pooling can route the next transaction to a
different
+ server connection, so a pooler with ``max_prepared_statements = 0`` fails
with
+ ``InvalidSqlStatementName`` once a statement outlives the backend that
prepared it,
+ or ``DuplicatePreparedStatement`` when a second client reuses the same
generated
+ name. PgBouncer supports prepared statements from 1.21 but only enables
them by
+ default from 1.24; Airflow's own chart pins 1.23.1 and sets
+ ``pool_mode = transaction``, so the recommended deployment is an affected
one.
+ This disables preparation outright, not just the automatic threshold, so
bulk
+ ``executemany`` writes re-parse per row rather than reusing one prepared
statement.
+ Overridable via ``[database] sql_alchemy_connect_args``.
+ """
+ if not sql_alchemy_conn or not sql_alchemy_conn.startswith(
Review Comment:
Would `make_url(sql_alchemy_conn).get_driver_name() in ("psycopg",
"psycopg_async")` be sturdier than the prefix match? An explicit
`sql_alchemy_conn_async` is used verbatim, and on SQLAlchemy 2.1 (allowed by
`>=2.0.50,!=2.1.0`) a bare `postgresql://` resolves to psycopg, so that config
would keep `prepare_threshold=5`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
