From 5aa9f394abcb0ca9637f74e6fe20b78bf82dc133 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Wed, 21 Oct 2026 18:38:39 +0000 Subject: [PATCH] feat(postgres): add pooler_mode, direct_dsn session split, and AlloyDB columnar GUC across asyncpg, psycopg, and psqlpy --- sqlspec/adapters/asyncpg/config.py | 63 ++++++++- sqlspec/adapters/asyncpg/core.py | 22 ++++ sqlspec/adapters/asyncpg/driver.py | 13 +- sqlspec/adapters/asyncpg/events/_hub.py | 47 ++++++- sqlspec/adapters/psqlpy/config.py | 61 ++++++++- sqlspec/adapters/psqlpy/core.py | 20 +++ sqlspec/adapters/psqlpy/driver.py | 34 ++++- sqlspec/adapters/psqlpy/events/_hub.py | 49 ++++++- sqlspec/adapters/psycopg/config.py | 121 +++++++++++++++++- sqlspec/adapters/psycopg/core.py | 8 ++ sqlspec/adapters/psycopg/driver.py | 24 ++++ sqlspec/adapters/psycopg/events/_hub.py | 92 ++++++++++++- sqlspec/migrations/base.py | 6 + sqlspec/migrations/commands.py | 4 +- .../adapters/_shared/_driver_type_system.py | 14 ++ .../unit/adapters/test_asyncpg/test_config.py | 107 +++++++++++++++- .../unit/adapters/test_psqlpy/test_config.py | 75 ++++++++++- .../unit/adapters/test_psycopg/test_config.py | 115 +++++++++++++++++ .../test_events/test_listener_hubs.py | 94 ++++++++++++++ .../migrations/test_migration_commands.py | 61 +++++++++ 20 files changed, 1005 insertions(+), 25 deletions(-) diff --git a/sqlspec/adapters/asyncpg/config.py b/sqlspec/adapters/asyncpg/config.py index 32423edd8..f17e4e115 100644 --- a/sqlspec/adapters/asyncpg/config.py +++ b/sqlspec/adapters/asyncpg/config.py @@ -1,8 +1,9 @@ """AsyncPG database configuration with direct field-based configuration.""" import asyncio +from collections.abc import Mapping from contextlib import suppress -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, TypedDict, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, Self, TypedDict, cast from mypy_extensions import mypyc_attr @@ -216,10 +217,14 @@ class AsyncpgDriverFeatures(TypedDict): alloydb_instance_uri: NotRequired[str] enable_alloydb_iam_auth: NotRequired[bool] alloydb_ip_type: NotRequired[str] + enable_alloydb_columnar: NotRequired[bool] enable_events: NotRequired[bool] events_backend: NotRequired[Literal["notify", "notify_queue", "poll_queue"]] on_connection_create: NotRequired["Callable[[AsyncpgConnection], Awaitable[None]]"] + pooler_mode: NotRequired[Literal["session", "transaction", "statement", "auto"] | None] pgbouncer: NotRequired[bool] + pgbouncer_mode: NotRequired[bool] + direct_dsn: NotRequired["str | Mapping[str, Any] | None"] type_codecs: NotRequired["list[dict[str, Any]]"] @@ -346,12 +351,30 @@ def __init__( self._user_connection_hook: Callable[[AsyncpgConnection], Awaitable[None]] | None = features_dict.get( "on_connection_create" ) - if connection_config and connection_config.get("pgbouncer"): + raw_conn_cfg = normalize_connection_config(connection_config) + self._explicit_statement_cache_size = "statement_cache_size" in raw_conn_cfg + raw_server_settings = raw_conn_cfg.get("server_settings") + self._explicit_server_settings: dict[str, str] | None = ( + dict(raw_server_settings) if isinstance(raw_server_settings, Mapping) else None + ) + normalized_conn_cfg = build_connection_config(raw_conn_cfg) + if raw_conn_cfg.get("pgbouncer"): features_dict["pgbouncer"] = True + features_dict["pgbouncer_mode"] = True + if features_dict.get("pooler_mode") in {None, "session", "auto"}: + features_dict["pooler_mode"] = "transaction" + if features_dict.get("pooler_mode") in {"transaction", "statement"} or bool( + features_dict.get("pgbouncer") or features_dict.get("pgbouncer_mode") + ): + normalized_conn_cfg.setdefault("statement_cache_size", 0) + if bool(features_dict.get("enable_alloydb_columnar", False)) and features_dict.get("pooler_mode") == "session": + merged_settings = dict(self._explicit_server_settings) if self._explicit_server_settings is not None else {} + merged_settings.setdefault("google_columnar_engine.enabled", "on") + normalized_conn_cfg["server_settings"] = merged_settings self._custom_type_codecs: list[dict[str, Any]] = list(features_dict.get("type_codecs") or []) super().__init__( - connection_config=build_connection_config(normalize_connection_config(connection_config)), + connection_config=normalized_conn_cfg, connection_instance=connection_instance, migration_config=migration_config, statement_config=statement_config, @@ -474,6 +497,40 @@ def get_event_runtime_hints(self) -> "EventRuntimeHints": return EventRuntimeHints(poll_interval=0.5, select_for_update=True, skip_locked=True) + def derive_direct_config(self) -> Self: + """Derive a session-mode configuration using direct_dsn for DDL, migrations, or LISTEN/NOTIFY.""" + direct_dsn = self.driver_features.get("direct_dsn") + if not direct_dsn: + return self + conn_cfg: dict[str, Any] = dict(self.connection_config) + if not self._explicit_statement_cache_size and conn_cfg.get("statement_cache_size") == 0: + conn_cfg.pop("statement_cache_size", None) + if self._explicit_server_settings is not None: + conn_cfg["server_settings"] = dict(self._explicit_server_settings) + else: + conn_cfg.pop("server_settings", None) + if isinstance(direct_dsn, str): + conn_cfg["dsn"] = direct_dsn + elif isinstance(direct_dsn, Mapping): + conn_cfg.update(build_connection_config(direct_dsn)) + features: dict[str, Any] = dict(self.driver_features) + features["pooler_mode"] = "session" + features["pgbouncer"] = False + features["pgbouncer_mode"] = False + features["direct_dsn"] = None + if self._user_connection_hook is not None: + features["on_connection_create"] = self._user_connection_hook + return type(self)( + connection_config=cast("AsyncpgPoolConfig", conn_cfg), + connection_instance=None, + migration_config=dict(self.migration_config), + statement_config=self.statement_config, + driver_features=cast("AsyncpgDriverFeatures", features), + bind_key=self.bind_key, + extension_config=cast("ExtensionConfigs", dict(self.extension_config)), + observability_config=self.observability_config, + ) + def _get_cloud_sql_connector(self) -> Any | None: """Return the configured Cloud SQL connector instance.""" return self._cloud_sql_connector diff --git a/sqlspec/adapters/asyncpg/core.py b/sqlspec/adapters/asyncpg/core.py index e3de29e89..384f59bb3 100644 --- a/sqlspec/adapters/asyncpg/core.py +++ b/sqlspec/adapters/asyncpg/core.py @@ -19,6 +19,7 @@ ConnectionTimeoutError, DeadlockError, ForeignKeyViolationError, + ImproperConfigurationError, IntegrityError, NotNullViolationError, OperationalError, @@ -280,6 +281,27 @@ def apply_driver_features( processed_features.setdefault("enable_alloydb", False) processed_features.setdefault("enable_events", True) processed_features.setdefault("events_backend", "notify") + pgbouncer_enabled = bool(processed_features.get("pgbouncer") or processed_features.get("pgbouncer_mode")) + processed_features["pgbouncer"] = pgbouncer_enabled + processed_features["pgbouncer_mode"] = pgbouncer_enabled + raw_pooler_mode = processed_features.get("pooler_mode") + if raw_pooler_mode is None: + processed_features["pooler_mode"] = "transaction" if pgbouncer_enabled else "session" + else: + normalized_mode = str(raw_pooler_mode).lower() + if normalized_mode not in {"session", "transaction", "statement", "auto"}: + msg = ( + f"Invalid pooler_mode: {raw_pooler_mode!r}. Expected 'session', 'transaction', 'statement', or 'auto'." + ) + raise ImproperConfigurationError(msg) + processed_features["pooler_mode"] = ( + "transaction" if normalized_mode == "auto" and pgbouncer_enabled else normalized_mode + ) + if processed_features["pooler_mode"] in {"transaction", "statement"}: + processed_features["pgbouncer"] = True + processed_features["pgbouncer_mode"] = True + processed_features.setdefault("direct_dsn", None) + processed_features.setdefault("enable_alloydb_columnar", False) parameter_config = configure_parameter_serializers( statement_config.parameter_config, serializer, deserializer=deserializer diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index 1fa49d816..ab2d27a63 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/sqlspec/adapters/asyncpg/driver.py @@ -236,6 +236,8 @@ async def begin(self) -> None: transaction = self.connection.transaction() try: await transaction.start() + if self._should_set_local_columnar(): + await self.connection.execute("SET LOCAL google_columnar_engine.enabled = 'on'") except AsyncpgPostgresError as e: self._release_failed_transaction_claim(transaction) raise create_mapped_exception(e) from e @@ -577,10 +579,19 @@ async def _run_stack_operations( async def _execute_stack_operation_prepared(self, normalized: "NormalizedStackOperation") -> StackResult: result = await self._execute_cached_statement(normalized.statement) - if not self.driver_features.get("pgbouncer"): + if not self._is_transaction_pooler_mode(): result.metadata["prepared_statement"] = True return StackResult.from_sql_result(result) + def _is_transaction_pooler_mode(self) -> bool: + pooler_mode = self.driver_features.get("pooler_mode", "session") + return pooler_mode in {"transaction", "statement"} or bool( + self.driver_features.get("pgbouncer") or self.driver_features.get("pgbouncer_mode") + ) + + def _should_set_local_columnar(self) -> bool: + return bool(self.driver_features.get("enable_alloydb_columnar", False)) and self._is_transaction_pooler_mode() + def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" return bool(self.connection.is_in_transaction()) diff --git a/sqlspec/adapters/asyncpg/events/_hub.py b/sqlspec/adapters/asyncpg/events/_hub.py index a9f99844e..5cafccc63 100644 --- a/sqlspec/adapters/asyncpg/events/_hub.py +++ b/sqlspec/adapters/asyncpg/events/_hub.py @@ -15,7 +15,7 @@ from typing import TYPE_CHECKING, Any, cast from weakref import WeakKeyDictionary -from sqlspec.exceptions import EventChannelError +from sqlspec.exceptions import EventChannelError, ImproperConfigurationError from sqlspec.extensions.events import normalize_event_channel_name from sqlspec.extensions.events._buffer import enqueue_with_capacity, resolve_listener_queue_capacity from sqlspec.utils.logging import get_logger, log_with_context @@ -40,6 +40,7 @@ class AsyncpgListenerHub: "_config", "_connection", "_connection_cm", + "_listener_config", "_listener_queue_capacity", "_lock", "_pool_destroying_registered", @@ -50,6 +51,7 @@ class AsyncpgListenerHub: def __init__(self, config: "AsyncpgConfig", backend_name: str = "notify") -> None: self._backend_name = backend_name self._config = config + self._listener_config: AsyncpgConfig | None = None self._lock = asyncio.Lock() self._listener_queue_capacity = resolve_listener_queue_capacity(config) self._queues: dict[str, WeakKeyDictionary[asyncio.Task[Any], asyncio.Queue[str]]] = {} @@ -143,6 +145,15 @@ async def shutdown(self) -> None: with contextlib.suppress(Exception): await connection_cm.__aexit__(None, None, None) _record_listener_lifecycle(self._config, self._backend_name, "release") + listener_config = self._listener_config + self._listener_config = None + if ( + listener_config is not None + and listener_config is not self._config + and callable(getattr(listener_config, "close_pool", None)) + ): + with contextlib.suppress(Exception): + await listener_config.close_pool() self._shutting_down = False def is_subscribed(self, channel: str) -> bool: @@ -162,6 +173,37 @@ def _get_consumer_queue(self, channel: str) -> "asyncio.Queue[str] | None": queues[task] = queue return queue + def _resolve_listener_config(self) -> "AsyncpgConfig": + if self._listener_config is not None: + return self._listener_config + raw_features = getattr(self._config, "driver_features", None) + driver_features = raw_features if isinstance(raw_features, dict) else {} + pooler_mode = str( + driver_features.get("pooler_mode") + or ( + "transaction" + if driver_features.get("pgbouncer") or driver_features.get("pgbouncer_mode") + else "session" + ) + ).lower() + raw_direct_dsn = driver_features.get("direct_dsn") + direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, str) and raw_direct_dsn else None + if pooler_mode in {"transaction", "statement"} and not direct_dsn: + logger.warning( + "PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode=%r without configuring driver_features['direct_dsn']", + pooler_mode, + ) + msg = ( + f"PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode={pooler_mode!r} without " + "configuring driver_features['direct_dsn']" + ) + raise ImproperConfigurationError(msg) + derive_fn = getattr(self._config, "derive_direct_config", None) + self._listener_config = ( + cast("AsyncpgConfig", derive_fn()) if direct_dsn and callable(derive_fn) else self._config + ) + return self._listener_config + async def _ensure_connection_locked(self) -> None: if self._connection is not None and not self._connection.is_closed(): return @@ -172,7 +214,8 @@ async def _ensure_connection_locked(self) -> None: _record_listener_lifecycle(self._config, self._backend_name, "release") self._connection = None self._connection_cm = None - connection_cm = self._config.provide_connection() + listener_config = self._resolve_listener_config() + connection_cm = listener_config.provide_connection() connection = await connection_cm.__aenter__() if connection is None or not has_add_listener(connection): with contextlib.suppress(Exception): diff --git a/sqlspec/adapters/psqlpy/config.py b/sqlspec/adapters/psqlpy/config.py index f04126f40..0abd69fc3 100644 --- a/sqlspec/adapters/psqlpy/config.py +++ b/sqlspec/adapters/psqlpy/config.py @@ -1,7 +1,8 @@ """Psqlpy database configuration.""" import sys -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, TypedDict, cast +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, Self, TypedDict, cast from mypy_extensions import mypyc_attr @@ -82,6 +83,7 @@ class PsqlpyPoolParams(PsqlpyConnectionParams): ports: NotRequired[list[int]] conn_recycling_method: NotRequired[str] max_db_pool_size: NotRequired[int] + max_cached_statement_lifetime: NotRequired[int] configure: NotRequired["Callable[..., Any]"] extra: NotRequired["dict[str, Any]"] @@ -127,6 +129,11 @@ class PsqlpyDriverFeatures(TypedDict): enable_pgvector: NotRequired[bool] enable_paradedb: NotRequired[bool] enable_pg_textsearch: NotRequired[bool] + enable_alloydb_columnar: NotRequired[bool] + pooler_mode: NotRequired[Literal["session", "transaction", "statement", "auto"] | None] + pgbouncer: NotRequired[bool] + pgbouncer_mode: NotRequired[bool] + direct_dsn: NotRequired["str | Mapping[str, Any] | None"] json_serializer: NotRequired["Callable[[Any], str]"] json_deserializer: NotRequired["Callable[[str], Any]"] on_connection_create: "NotRequired[Callable[[PsqlpyConnection], Awaitable[None]]]" @@ -245,12 +252,28 @@ def __init__( extension_config: Extension-specific configuration. **kwargs: Additional keyword arguments. """ - connection_config = build_connection_config(normalize_connection_config(connection_config)) + raw_conn_cfg = normalize_connection_config(connection_config) + self._explicit_max_cached_statement_lifetime = "max_cached_statement_lifetime" in raw_conn_cfg + raw_options = raw_conn_cfg.get("options") + self._explicit_options: str | None = raw_options if isinstance(raw_options, str) else None + connection_config = build_connection_config(raw_conn_cfg) statement_config = statement_config or default_statement_config statement_config, driver_features = apply_driver_features(statement_config, driver_features) features_dict = dict(driver_features) if driver_features else {} + if features_dict.get("pooler_mode") in {"transaction", "statement"} or bool( + features_dict.get("pgbouncer") or features_dict.get("pgbouncer_mode") + ): + connection_config.setdefault("max_cached_statement_lifetime", 0) + if bool(features_dict.get("enable_alloydb_columnar", False)) and features_dict.get("pooler_mode") == "session": + columnar_opt = "-c google_columnar_engine.enabled=on" + existing_options = connection_config.get("options") + if isinstance(existing_options, str) and existing_options.strip(): + if "google_columnar_engine.enabled" not in existing_options: + connection_config["options"] = f"{existing_options} {columnar_opt}" + else: + connection_config["options"] = columnar_opt self._user_connection_hook: Callable[[PsqlpyConnection], Awaitable[None]] | None = features_dict.get( "on_connection_create" ) @@ -349,6 +372,40 @@ def get_event_runtime_hints(self) -> "EventRuntimeHints": return EventRuntimeHints(poll_interval=0.5, select_for_update=True, skip_locked=True) + def derive_direct_config(self) -> Self: + """Derive a session-mode configuration using direct_dsn for DDL, migrations, or LISTEN/NOTIFY.""" + direct_dsn = self.driver_features.get("direct_dsn") + if not direct_dsn: + return self + conn_cfg: dict[str, Any] = dict(self.connection_config) + if not self._explicit_max_cached_statement_lifetime and conn_cfg.get("max_cached_statement_lifetime") == 0: + conn_cfg.pop("max_cached_statement_lifetime", None) + if self._explicit_options is not None: + conn_cfg["options"] = self._explicit_options + else: + conn_cfg.pop("options", None) + if isinstance(direct_dsn, str): + conn_cfg["dsn"] = direct_dsn + elif isinstance(direct_dsn, Mapping): + conn_cfg.update(build_connection_config(direct_dsn)) + features: dict[str, Any] = dict(self.driver_features) + features["pooler_mode"] = "session" + features["pgbouncer"] = False + features["pgbouncer_mode"] = False + features["direct_dsn"] = None + if self._user_connection_hook is not None: + features["on_connection_create"] = self._user_connection_hook + return type(self)( + connection_config=cast("PsqlpyPoolParams", conn_cfg), + connection_instance=None, + migration_config=dict(self.migration_config), + statement_config=self.statement_config, + driver_features=cast("PsqlpyDriverFeatures", features), + bind_key=self.bind_key, + extension_config=cast("ExtensionConfigs", dict(self.extension_config)), + observability_config=self.observability_config, + ) + async def _ensure_connection(self, connection: "PsqlpyConnection") -> None: """Ensure connection callback has been called once for this connection. diff --git a/sqlspec/adapters/psqlpy/core.py b/sqlspec/adapters/psqlpy/core.py index 05e4b5c4f..efd31dc3a 100644 --- a/sqlspec/adapters/psqlpy/core.py +++ b/sqlspec/adapters/psqlpy/core.py @@ -37,6 +37,7 @@ DataError, DeadlockError, ForeignKeyViolationError, + ImproperConfigurationError, IntegrityError, NotNullViolationError, OperationalError, @@ -183,6 +184,25 @@ def apply_driver_features( features.setdefault("enable_pgvector", PGVECTOR_INSTALLED) features.setdefault("enable_paradedb", True) features.setdefault("enable_pg_textsearch", True) + pgbouncer_enabled = bool(features.get("pgbouncer") or features.get("pgbouncer_mode")) + features["pgbouncer"] = pgbouncer_enabled + features["pgbouncer_mode"] = pgbouncer_enabled + raw_pooler_mode = features.get("pooler_mode") + if raw_pooler_mode is None: + features["pooler_mode"] = "transaction" if pgbouncer_enabled else "session" + else: + normalized_mode = str(raw_pooler_mode).lower() + if normalized_mode not in {"session", "transaction", "statement", "auto"}: + msg = ( + f"Invalid pooler_mode: {raw_pooler_mode!r}. Expected 'session', 'transaction', 'statement', or 'auto'." + ) + raise ImproperConfigurationError(msg) + features["pooler_mode"] = "transaction" if normalized_mode == "auto" and pgbouncer_enabled else normalized_mode + if features["pooler_mode"] in {"transaction", "statement"}: + features["pgbouncer"] = True + features["pgbouncer_mode"] = True + features.setdefault("direct_dsn", None) + features.setdefault("enable_alloydb_columnar", False) base_config = build_statement_config_from_profile( driver_profile, json_serializer=serializer, json_deserializer=deserializer diff --git a/sqlspec/adapters/psqlpy/driver.py b/sqlspec/adapters/psqlpy/driver.py index ce015f55d..f5b722c5c 100644 --- a/sqlspec/adapters/psqlpy/driver.py +++ b/sqlspec/adapters/psqlpy/driver.py @@ -123,7 +123,11 @@ async def dispatch_execute(self, cursor: "PsqlpyConnection", statement: SQL) -> params = cast("Sequence[Any] | Mapping[str, Any] | None", prepared_parameters) or [] if statement.returns_rows(): - query_result = await cursor.fetch(sql, params) + query_result = ( + await cursor.fetch(sql, params, prepared=False) + if self._is_transaction_pooler_mode() + else await cursor.fetch(sql, params) + ) dict_rows, column_names = collect_rows(query_result) return self.create_execution_result( @@ -138,7 +142,11 @@ async def dispatch_execute(self, cursor: "PsqlpyConnection", statement: SQL) -> if statement.operation_type in {"INSERT", "UPDATE", "DELETE"}: count_sql = _dml_count_query(sql) if count_sql is not None: - count_result = await cursor.fetch(count_sql, params) + count_result = ( + await cursor.fetch(count_sql, params, prepared=False) + if self._is_transaction_pooler_mode() + else await cursor.fetch(count_sql, params) + ) count_rows, _ = collect_rows(count_result) if len(count_rows) != 1 or set(count_rows[0]) != {_DML_COUNT_COLUMN}: msg = "psqlpy DML row count query returned an invalid result" @@ -149,7 +157,11 @@ async def dispatch_execute(self, cursor: "PsqlpyConnection", statement: SQL) -> raise SQLSpecError(msg) return self.create_execution_result(cursor, rowcount_override=rows_affected) - result = await cursor.execute(sql, params) + result = ( + await cursor.execute(sql, params, prepared=False) + if self._is_transaction_pooler_mode() + else await cursor.execute(sql, params) + ) rows_affected = extract_rows_affected(result) return self.create_execution_result(cursor, rowcount_override=rows_affected) @@ -171,7 +183,10 @@ async def dispatch_execute_many(self, cursor: "PsqlpyConnection", statement: SQL formatted_parameters = format_execute_many_parameters(prepared_parameters, coerce_numeric=False) - await cursor.execute_many(sql, formatted_parameters) + if self._is_transaction_pooler_mode(): + await cursor.execute_many(sql, formatted_parameters, prepared=False) + else: + await cursor.execute_many(sql, formatted_parameters) rows_affected = len(formatted_parameters) @@ -213,6 +228,8 @@ async def begin(self) -> None: try: await self.connection.execute("BEGIN") + if self._should_set_local_columnar(): + await self.connection.execute("SET LOCAL google_columnar_engine.enabled = 'on'") except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError) as e: msg = f"Failed to begin psqlpy transaction: {e}" raise SQLSpecError(msg) from e @@ -488,5 +505,14 @@ def _connection_in_transaction(self) -> bool: """ return self._transaction_active + def _is_transaction_pooler_mode(self) -> bool: + pooler_mode = self.driver_features.get("pooler_mode", "session") + return pooler_mode in {"transaction", "statement"} or bool( + self.driver_features.get("pgbouncer") or self.driver_features.get("pgbouncer_mode") + ) + + def _should_set_local_columnar(self) -> bool: + return bool(self.driver_features.get("enable_alloydb_columnar", False)) and self._is_transaction_pooler_mode() + register_driver_profile("psqlpy", driver_profile) diff --git a/sqlspec/adapters/psqlpy/events/_hub.py b/sqlspec/adapters/psqlpy/events/_hub.py index 86ff29853..20da69998 100644 --- a/sqlspec/adapters/psqlpy/events/_hub.py +++ b/sqlspec/adapters/psqlpy/events/_hub.py @@ -17,6 +17,7 @@ from weakref import WeakKeyDictionary from sqlspec.core import SQL +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.events import normalize_event_channel_name from sqlspec.extensions.events._buffer import enqueue_with_capacity, resolve_listener_queue_capacity from sqlspec.utils.logging import get_logger, log_with_context @@ -80,6 +81,7 @@ class PsqlpyListenerHub: "_callbacks", "_config", "_listener", + "_listener_config", "_listener_queue_capacity", "_listener_started", "_lock", @@ -92,6 +94,7 @@ class PsqlpyListenerHub: def __init__(self, config: "PsqlpyConfig", backend_name: str = "notify") -> None: self._backend_name = backend_name self._config = config + self._listener_config: PsqlpyConfig | None = None self._lock = asyncio.Lock() self._listener_queue_capacity = resolve_listener_queue_capacity(config) self._queues: dict[str, WeakKeyDictionary[asyncio.Task[Any], asyncio.Queue[str]]] = {} @@ -168,6 +171,15 @@ async def shutdown(self) -> None: with contextlib.suppress(Exception): await listener.shutdown() _record_listener_lifecycle(self._config, self._backend_name, "release") + listener_config = self._listener_config + self._listener_config = None + if ( + listener_config is not None + and listener_config is not self._config + and callable(getattr(listener_config, "close_pool", None)) + ): + with contextlib.suppress(Exception): + await listener_config.close_pool() self._shutting_down = False def is_subscribed(self, channel: str) -> bool: @@ -220,6 +232,37 @@ async def _subscribe_locked( raise return consumer_queue + def _resolve_listener_config(self) -> "PsqlpyConfig": + if self._listener_config is not None: + return self._listener_config + raw_features = getattr(self._config, "driver_features", None) + driver_features = raw_features if isinstance(raw_features, dict) else {} + pooler_mode = str( + driver_features.get("pooler_mode") + or ( + "transaction" + if driver_features.get("pgbouncer") or driver_features.get("pgbouncer_mode") + else "session" + ) + ).lower() + raw_direct_dsn = driver_features.get("direct_dsn") + direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, str) and raw_direct_dsn else None + if pooler_mode in {"transaction", "statement"} and not direct_dsn: + logger.warning( + "PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode=%r without configuring driver_features['direct_dsn']", + pooler_mode, + ) + msg = ( + f"PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode={pooler_mode!r} without " + "configuring driver_features['direct_dsn']" + ) + raise ImproperConfigurationError(msg) + derive_fn = getattr(self._config, "derive_direct_config", None) + self._listener_config = ( + cast("PsqlpyConfig", derive_fn()) if direct_dsn and callable(derive_fn) else self._config + ) + return self._listener_config + async def _ensure_listener_locked(self) -> "PsqlpyListener": if self._listener is not None and self._listener.is_started: return cast("PsqlpyListener", self._listener) @@ -234,7 +277,8 @@ async def _ensure_listener_locked(self) -> "PsqlpyListener": _record_listener_lifecycle(self._config, self._backend_name, "release") self._listener = None self._listener_started = False - pool = await self._config.provide_pool() + listener_config = self._resolve_listener_config() + pool = await listener_config.provide_pool() listener = pool.listener() await listener.startup() self._listener = listener @@ -271,7 +315,8 @@ async def _wait_until_ready(self, channel: str, *, timeout: float) -> None: ready = asyncio.Event() self._ready_events.setdefault(channel, {})[payload] = ready try: - async with self._config.provide_session() as driver: + listener_config = self._resolve_listener_config() + async with listener_config.provide_session() as driver: await driver.execute_script(SQL("SELECT pg_notify($1, $2)", channel, payload)) await driver.commit() await asyncio.wait_for(ready.wait(), timeout=timeout) diff --git a/sqlspec/adapters/psycopg/config.py b/sqlspec/adapters/psycopg/config.py index c7a39db80..46437bfff 100644 --- a/sqlspec/adapters/psycopg/config.py +++ b/sqlspec/adapters/psycopg/config.py @@ -100,6 +100,7 @@ "timeout", }) _ALLOYDB_CONNECTOR_PACKAGE = "google-cloud-alloydb-connector" +_DEFAULT_PSYCOPG_PREPARE_THRESHOLD: Final[int] = 5 class PsycopgConnectionParams(TypedDict): @@ -191,6 +192,18 @@ class PsycopgDriverFeatures(TypedDict): enable_alloydb_iam_auth: Enable AlloyDB IAM database authentication for sync connector connections. Defaults to False. alloydb_ip_type: AlloyDB connector IP type. Defaults to PRIVATE. + enable_alloydb_columnar: Enable Google AlloyDB Columnar Engine acceleration. + Configures ``google_columnar_engine.enabled=on`` at the session level via ``options`` + in session mode, or per-transaction via ``SET LOCAL`` in transaction/statement pooler mode. + Defaults to False. + pooler_mode: Connection pooler mode ("session", "transaction", "statement", "auto"). + When "transaction" or "statement", sets ``prepare_threshold=None`` to disable server-side + prepared statements for PgBouncer / PgCat / Odyssey / Supavisor compatibility. + Defaults to "session". + pgbouncer: Legacy alias for ``pooler_mode="transaction"``. + pgbouncer_mode: Legacy alias for ``pooler_mode="transaction"``. + direct_dsn: Optional direct PostgreSQL DSN or connection config mapping bypassing the external + pooler for session-scoped operations (DDL/migrations and LISTEN/NOTIFY). null_pool: Enable NullConnectionPool / AsyncNullConnectionPool for serverless / PgBouncer environments. """ @@ -206,6 +219,11 @@ class PsycopgDriverFeatures(TypedDict): alloydb_instance_uri: NotRequired[str] enable_alloydb_iam_auth: NotRequired[bool] alloydb_ip_type: NotRequired[str] + enable_alloydb_columnar: NotRequired[bool] + pooler_mode: NotRequired[Literal["session", "transaction", "statement", "auto"]] + pgbouncer: NotRequired[bool] + pgbouncer_mode: NotRequired[bool] + direct_dsn: NotRequired["str | Mapping[str, Any] | None"] null_pool: NotRequired[bool] @@ -329,6 +347,21 @@ def __init__( features_dict.setdefault("enable_alloydb", False) features_dict.setdefault("enable_alloydb_iam_auth", False) features_dict.setdefault("alloydb_ip_type", "PRIVATE") + pooler_mode = features_dict.get("pooler_mode", "session") + is_transaction_pooler = bool( + pooler_mode in {"transaction", "statement"} + or features_dict.get("pgbouncer") + or features_dict.get("pgbouncer_mode") + ) + if is_transaction_pooler: + connection_config.setdefault("prepare_threshold", None) + if features_dict.get("enable_alloydb_columnar") and not is_transaction_pooler: + existing_options = connection_config.get("options") + columnar_opt = "-c google_columnar_engine.enabled=on" + if not existing_options: + connection_config["options"] = columnar_opt + elif columnar_opt not in str(existing_options): + connection_config["options"] = f"{existing_options} {columnar_opt}" self._pgvector_available: bool | None = None self._paradedb_available: bool | None = None self._pg_textsearch_available: bool | None = None @@ -348,6 +381,37 @@ def __init__( self._validate_alloydb_config() + def derive_direct_config(self) -> Self: + """Derive a session-mode Psycopg sync configuration targeting direct_dsn when configured.""" + direct_dsn = self.driver_features.get("direct_dsn") + if not direct_dsn: + return self + conn_cfg = dict(self.connection_config) + if isinstance(direct_dsn, Mapping): + conn_cfg.update(build_connection_config(direct_dsn)) + else: + conn_cfg["conninfo"] = str(direct_dsn) + is_pooled = bool( + self.driver_features.get("pooler_mode") in {"transaction", "statement"} + or self.driver_features.get("pgbouncer") + or self.driver_features.get("pgbouncer_mode") + ) + if is_pooled and conn_cfg.get("prepare_threshold") is None: + conn_cfg.pop("prepare_threshold", None) + features = dict(self.driver_features) + features["pooler_mode"] = "session" + features["pgbouncer"] = False + features["pgbouncer_mode"] = False + features["direct_dsn"] = None + return type(self)( + connection_config=conn_cfg, + migration_config=dict(self.migration_config) if self.migration_config else None, + statement_config=self.statement_config, + driver_features=features, + bind_key=self.bind_key, + extension_config=dict(self.extension_config) if self.extension_config else None, + ) + @property def pg_textsearch_available(self) -> bool: """Return True if the pg_textsearch extension is available.""" @@ -469,6 +533,7 @@ def _setup_alloydb_connector( user = config.get("user") password = config.get("password") database = config.get("dbname") or config.get("db") + prepare_threshold = config.pop("prepare_threshold", _DEFAULT_PSYCOPG_PREPARE_THRESHOLD) for key in _ALLOYDB_DIRECT_CONNECTION_KEYS: config.pop(key, None) @@ -482,6 +547,7 @@ def _setup_alloydb_connector( user=cast("str | None", user), password=cast("str | None", password), database=cast("str | None", database), + prepare_threshold=cast("int | None", prepare_threshold), ) def _create_pool(self) -> "ConnectionPool": @@ -729,6 +795,21 @@ def __init__( self._user_connection_hook: Callable[[PsycopgAsyncConnection], Awaitable[None]] | None = features_dict.get( "on_connection_create" ) + pooler_mode = features_dict.get("pooler_mode", "session") + is_transaction_pooler = bool( + pooler_mode in {"transaction", "statement"} + or features_dict.get("pgbouncer") + or features_dict.get("pgbouncer_mode") + ) + if is_transaction_pooler: + connection_config.setdefault("prepare_threshold", None) + if features_dict.get("enable_alloydb_columnar") and not is_transaction_pooler: + existing_options = connection_config.get("options") + columnar_opt = "-c google_columnar_engine.enabled=on" + if not existing_options: + connection_config["options"] = columnar_opt + elif columnar_opt not in str(existing_options): + connection_config["options"] = f"{existing_options} {columnar_opt}" self._pgvector_available: bool | None = None self._paradedb_available: bool | None = None self._pg_textsearch_available: bool | None = None @@ -745,6 +826,37 @@ def __init__( **kwargs, ) + def derive_direct_config(self) -> Self: + """Derive a session-mode Psycopg async configuration targeting direct_dsn when configured.""" + direct_dsn = self.driver_features.get("direct_dsn") + if not direct_dsn: + return self + conn_cfg = dict(self.connection_config) + if isinstance(direct_dsn, Mapping): + conn_cfg.update(build_connection_config(direct_dsn)) + else: + conn_cfg["conninfo"] = str(direct_dsn) + is_pooled = bool( + self.driver_features.get("pooler_mode") in {"transaction", "statement"} + or self.driver_features.get("pgbouncer") + or self.driver_features.get("pgbouncer_mode") + ) + if is_pooled and conn_cfg.get("prepare_threshold") is None: + conn_cfg.pop("prepare_threshold", None) + features = dict(self.driver_features) + features["pooler_mode"] = "session" + features["pgbouncer"] = False + features["pgbouncer_mode"] = False + features["direct_dsn"] = None + return type(self)( + connection_config=conn_cfg, + migration_config=dict(self.migration_config) if self.migration_config else None, + statement_config=self.statement_config, + driver_features=features, + bind_key=self.bind_key, + extension_config=dict(self.extension_config) if self.extension_config else None, + ) + @property def pg_textsearch_available(self) -> bool: """Return True if the pg_textsearch extension is available.""" @@ -970,19 +1082,20 @@ def _make_alloydb_connection_class( user: str | None, password: str | None, database: str | None, + prepare_threshold: int | None = _DEFAULT_PSYCOPG_PREPARE_THRESHOLD, ) -> "type[PsycopgConnection[Any]]": """Build a psycopg connection class backed by the AlloyDB connector.""" - default_prepare_threshold = 5 - class _AlloyDBPsycopgConnection(PsycopgConnection[Any]): + default_prepare_threshold: ClassVar[int | None] = prepare_threshold + @classmethod def connect( cls, conninfo: str = "", *, autocommit: bool = False, - prepare_threshold: int | None = default_prepare_threshold, + prepare_threshold: int | None = prepare_threshold, context: "AdaptContext | None" = None, row_factory: "RowFactory[Any] | None" = None, cursor_factory: "type[Cursor[Any]] | None" = None, @@ -1001,7 +1114,7 @@ def connect( conn = connector.connect(instance_uri, "psycopg", **connector_kwargs) if autocommit: conn.autocommit = autocommit - if prepare_threshold != default_prepare_threshold: + if prepare_threshold != _DEFAULT_PSYCOPG_PREPARE_THRESHOLD: conn.prepare_threshold = prepare_threshold if row_factory is not None: conn.row_factory = row_factory diff --git a/sqlspec/adapters/psycopg/core.py b/sqlspec/adapters/psycopg/core.py index ee9010839..9653411cf 100644 --- a/sqlspec/adapters/psycopg/core.py +++ b/sqlspec/adapters/psycopg/core.py @@ -177,6 +177,14 @@ def apply_driver_features( features.setdefault("enable_pgvector", PGVECTOR_INSTALLED) features.setdefault("enable_paradedb", True) features.setdefault("enable_pg_textsearch", True) + features.setdefault("enable_alloydb_columnar", False) + raw_pooler_mode = features.get("pooler_mode") + if raw_pooler_mode is None and (features.get("pgbouncer") or features.get("pgbouncer_mode")): + raw_pooler_mode = "transaction" + features["pooler_mode"] = raw_pooler_mode or "session" + features.setdefault("pgbouncer", False) + features.setdefault("pgbouncer_mode", False) + features.setdefault("direct_dsn", None) parameter_config = _parameter_config(driver_profile, serializer, deserializer) statement_config = statement_config.replace(parameter_config=parameter_config) diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index d501f55d2..8e2229c4d 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -391,6 +391,8 @@ def begin(self) -> None: restore_autocommit = self.connection.autocommit if restore_autocommit: self.connection.autocommit = False + if self._should_set_local_columnar(): + self.connection.execute("SET LOCAL google_columnar_engine.enabled = 'on'") except psycopg.Error as e: mapped = create_mapped_exception(e) if type(mapped) is not SQLSpecError: @@ -689,6 +691,16 @@ def _restore_original_autocommit(self) -> None: msg = f"Failed to restore autocommit: {e}" raise SQLSpecError(msg) from e + def _should_set_local_columnar(self) -> bool: + """Return True when AlloyDB columnar GUC must be set transaction-locally.""" + if not self.driver_features.get("enable_alloydb_columnar"): + return False + return bool( + self.driver_features.get("pooler_mode") in {"transaction", "statement"} + or self.driver_features.get("pgbouncer") + or self.driver_features.get("pgbouncer_mode") + ) + class PsycopgAsyncExceptionHandler(BaseAsyncExceptionHandler): """Async context manager for handling PostgreSQL psycopg database exceptions. @@ -910,6 +922,8 @@ async def begin(self) -> None: restore_autocommit = self.connection.autocommit if restore_autocommit: await self.connection.set_autocommit(False) + if self._should_set_local_columnar(): + await self.connection.execute("SET LOCAL google_columnar_engine.enabled = 'on'") except psycopg.Error as e: mapped = create_mapped_exception(e) if type(mapped) is not SQLSpecError: @@ -1215,6 +1229,16 @@ async def _restore_original_autocommit(self) -> None: msg = f"Failed to restore autocommit: {e}" raise SQLSpecError(msg) from e + def _should_set_local_columnar(self) -> bool: + """Return True when AlloyDB columnar GUC must be set transaction-locally.""" + if not self.driver_features.get("enable_alloydb_columnar"): + return False + return bool( + self.driver_features.get("pooler_mode") in {"transaction", "statement"} + or self.driver_features.get("pgbouncer") + or self.driver_features.get("pgbouncer_mode") + ) + def _attribute_pipeline_failure( pending: "list[PipelineCursorEntry]", current: "PreparedStackOperation" diff --git a/sqlspec/adapters/psycopg/events/_hub.py b/sqlspec/adapters/psycopg/events/_hub.py index a9a98d5ab..d0634504e 100644 --- a/sqlspec/adapters/psycopg/events/_hub.py +++ b/sqlspec/adapters/psycopg/events/_hub.py @@ -14,7 +14,7 @@ import logging import queue as stdlib_queue import threading -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from weakref import WeakKeyDictionary from sqlspec.adapters.psycopg._typing import PsycopgComposed, PsycopgIdentifier, PsycopgSQL @@ -44,6 +44,7 @@ class PsycopgAsyncListenerHub: "_config", "_connection", "_connection_cm", + "_listener_config", "_listener_queue_capacity", "_lock", "_pool_destroying_registered", @@ -56,6 +57,7 @@ class PsycopgAsyncListenerHub: def __init__(self, config: "PsycopgAsyncConfig", backend_name: str = "notify") -> None: self._backend_name = backend_name self._config = config + self._listener_config: PsycopgAsyncConfig | None = None self._lock = asyncio.Lock() self._listener_queue_capacity = resolve_listener_queue_capacity(config) self._queues: dict[str, WeakKeyDictionary[asyncio.Task[Any], asyncio.Queue[str]]] = {} @@ -137,6 +139,15 @@ async def shutdown(self) -> None: with contextlib.suppress(Exception): await connection_cm.__aexit__(None, None, None) _record_listener_lifecycle(self._config, self._backend_name, "release", is_async=True) + listener_config = self._listener_config + self._listener_config = None + if ( + listener_config is not None + and listener_config is not self._config + and callable(getattr(listener_config, "close_pool", None)) + ): + with contextlib.suppress(Exception): + await listener_config.close_pool() self._shutting_down = False self._stopping = False @@ -157,6 +168,37 @@ def _get_consumer_queue(self, channel: str) -> "asyncio.Queue[str] | None": queues[task] = queue return queue + def _resolve_listener_config(self) -> "PsycopgAsyncConfig": + if self._listener_config is not None: + return self._listener_config + raw_features = getattr(self._config, "driver_features", None) + driver_features = raw_features if isinstance(raw_features, dict) else {} + pooler_mode = str( + driver_features.get("pooler_mode") + or ( + "transaction" + if driver_features.get("pgbouncer") or driver_features.get("pgbouncer_mode") + else "session" + ) + ).lower() + raw_direct_dsn = driver_features.get("direct_dsn") + direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, str) and raw_direct_dsn else None + if pooler_mode in {"transaction", "statement"} and not direct_dsn: + logger.warning( + "PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode=%r without configuring driver_features['direct_dsn']", + pooler_mode, + ) + msg = ( + f"PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode={pooler_mode!r} without " + "configuring driver_features['direct_dsn']" + ) + raise ImproperConfigurationError(msg) + derive_fn = getattr(self._config, "derive_direct_config", None) + self._listener_config = ( + cast("PsycopgAsyncConfig", derive_fn()) if direct_dsn and callable(derive_fn) else self._config + ) + return self._listener_config + async def _ensure_connection_locked(self) -> None: if ( self._connection is not None @@ -179,7 +221,8 @@ async def _ensure_connection_locked(self) -> None: self._connection_cm = None self._pump_task = None self._stopping = False - connection_cm = self._config.provide_connection() + listener_config = self._resolve_listener_config() + connection_cm = listener_config.provide_connection() connection = await connection_cm.__aenter__() if connection is None: with contextlib.suppress(Exception): @@ -273,6 +316,7 @@ class PsycopgSyncListenerHub: "_config", "_connection", "_connection_cm", + "_listener_config", "_listener_queue_capacity", "_lock", "_pool_destroying_registered", @@ -285,6 +329,7 @@ class PsycopgSyncListenerHub: def __init__(self, config: "PsycopgSyncConfig", backend_name: str = "notify") -> None: self._backend_name = backend_name self._config = config + self._listener_config: PsycopgSyncConfig | None = None self._lock = threading.Lock() self._listener_queue_capacity = resolve_listener_queue_capacity(config) self._queues: dict[str, WeakKeyDictionary[threading.Thread, stdlib_queue.Queue[str]]] = {} @@ -361,6 +406,15 @@ def shutdown(self) -> None: with contextlib.suppress(Exception): connection_cm.__exit__(None, None, None) _record_listener_lifecycle(self._config, self._backend_name, "release", is_async=False) + listener_config = self._listener_config + self._listener_config = None + if ( + listener_config is not None + and listener_config is not self._config + and callable(getattr(listener_config, "close_pool", None)) + ): + with contextlib.suppress(Exception): + listener_config.close_pool() self._connection = None self._stopping = threading.Event() self._command_queue = stdlib_queue.Queue() @@ -381,6 +435,37 @@ def _get_consumer_queue(self, channel: str) -> "stdlib_queue.Queue[str] | None": queues[thread] = queue return queue + def _resolve_listener_config(self) -> "PsycopgSyncConfig": + if self._listener_config is not None: + return self._listener_config + raw_features = getattr(self._config, "driver_features", None) + driver_features = raw_features if isinstance(raw_features, dict) else {} + pooler_mode = str( + driver_features.get("pooler_mode") + or ( + "transaction" + if driver_features.get("pgbouncer") or driver_features.get("pgbouncer_mode") + else "session" + ) + ).lower() + raw_direct_dsn = driver_features.get("direct_dsn") + direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, str) and raw_direct_dsn else None + if pooler_mode in {"transaction", "statement"} and not direct_dsn: + logger.warning( + "PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode=%r without configuring driver_features['direct_dsn']", + pooler_mode, + ) + msg = ( + f"PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode={pooler_mode!r} without " + "configuring driver_features['direct_dsn']" + ) + raise ImproperConfigurationError(msg) + derive_fn = getattr(self._config, "derive_direct_config", None) + self._listener_config = ( + cast("PsycopgSyncConfig", derive_fn()) if direct_dsn and callable(derive_fn) else self._config + ) + return self._listener_config + def _ensure_connection_locked(self) -> None: if ( self._connection is not None @@ -405,7 +490,8 @@ def _ensure_connection_locked(self) -> None: self._worker_thread = None self._stopping = threading.Event() self._command_queue = stdlib_queue.Queue() - connection_cm = self._config.provide_connection() + listener_config = self._resolve_listener_config() + connection_cm = listener_config.provide_connection() connection = connection_cm.__enter__() if connection is None: with contextlib.suppress(Exception): diff --git a/sqlspec/migrations/base.py b/sqlspec/migrations/base.py index 3fdb77bc0..e16f54dc9 100644 --- a/sqlspec/migrations/base.py +++ b/sqlspec/migrations/base.py @@ -447,6 +447,12 @@ def __init__(self, config: ConfigT) -> None: Args: config: The SQLSpec configuration. """ + raw_features = getattr(config, "driver_features", None) + driver_features = raw_features if isinstance(raw_features, dict) else {} + direct_dsn = driver_features.get("direct_dsn") + derive_fn = getattr(config, "derive_direct_config", None) + if isinstance(direct_dsn, str) and direct_dsn and callable(derive_fn): + config = cast("ConfigT", derive_fn()) self.config = config migration_config = self._get_migration_config() diff --git a/sqlspec/migrations/commands.py b/sqlspec/migrations/commands.py index e77287744..75462af4b 100644 --- a/sqlspec/migrations/commands.py +++ b/sqlspec/migrations/commands.py @@ -161,7 +161,7 @@ def __init__(self, config: "SyncConfigT") -> None: """ super().__init__(config) self.tracker = self._create_tracker() - context = MigrationContext.from_config(config) + context = MigrationContext.from_config(self.config) context.extension_config = self.extension_configs self.runner = SyncMigrationRunner( @@ -1122,7 +1122,7 @@ def __init__(self, config: "AsyncConfigT") -> None: """ super().__init__(config) self.tracker = self._create_tracker() - context = MigrationContext.from_config(config) + context = MigrationContext.from_config(self.config) context.extension_config = self.extension_configs self.runner = AsyncMigrationRunner( diff --git a/tests/integration/adapters/_shared/_driver_type_system.py b/tests/integration/adapters/_shared/_driver_type_system.py index 629e15b55..7601c431d 100644 --- a/tests/integration/adapters/_shared/_driver_type_system.py +++ b/tests/integration/adapters/_shared/_driver_type_system.py @@ -178,10 +178,14 @@ class SourceEquivalenceCase: "alloydb_instance_uri", "enable_alloydb_iam_auth", "alloydb_ip_type", + "enable_alloydb_columnar", "enable_events", "events_backend", "on_connection_create", + "pooler_mode", "pgbouncer", + "pgbouncer_mode", + "direct_dsn", "type_codecs", ), "bigquery": ( @@ -292,11 +296,16 @@ class SourceEquivalenceCase: "enable_pgvector", "enable_paradedb", "enable_pg_textsearch", + "enable_alloydb_columnar", "json_serializer", "json_deserializer", "on_connection_create", "enable_events", "events_backend", + "pooler_mode", + "pgbouncer", + "pgbouncer_mode", + "direct_dsn", ), "psycopg": ( "enable_pgvector", @@ -311,6 +320,11 @@ class SourceEquivalenceCase: "alloydb_instance_uri", "enable_alloydb_iam_auth", "alloydb_ip_type", + "enable_alloydb_columnar", + "pooler_mode", + "pgbouncer", + "pgbouncer_mode", + "direct_dsn", "null_pool", ), "pymssql": ("json_serializer", "json_deserializer", "on_connection_create", "enable_events", "events_backend"), diff --git a/tests/unit/adapters/test_asyncpg/test_config.py b/tests/unit/adapters/test_asyncpg/test_config.py index ec00b7cab..35e772ed3 100644 --- a/tests/unit/adapters/test_asyncpg/test_config.py +++ b/tests/unit/adapters/test_asyncpg/test_config.py @@ -20,7 +20,8 @@ build_statement_config, resolve_postgres_extension_state, ) -from sqlspec.core import StatementConfig +from sqlspec.adapters.asyncpg.driver import AsyncpgDriver +from sqlspec.core import StatementConfig, StatementStack from sqlspec.exceptions import ImproperConfigurationError @@ -389,3 +390,107 @@ async def test_asyncpg_extension_probe_retries_after_transient_error() -> None: assert config._pgvector_available is False # pyright: ignore[reportPrivateUsage] assert config.pg_textsearch_available is True assert config.statement_config.dialect == "pg_textsearch" + + +async def test_asyncpg_pooler_mode_transaction_sets_statement_cache_size_zero(monkeypatch: pytest.MonkeyPatch) -> None: + """Transaction pooler mode and pgbouncer aliases should disable statement_cache_size and prepared stack metadata.""" + captured_connect_kwargs: dict[str, Any] = {} + + async def fake_create_pool(*, init: Any = None, **kwargs: Any) -> object: + _ = init + captured_connect_kwargs.update(kwargs) + return object() + + monkeypatch.setattr(config_mod, "asyncpg_create_pool", fake_create_pool) + + config = AsyncpgConfig( + connection_config={"dsn": "postgresql://localhost/test"}, driver_features={"pooler_mode": "transaction"} + ) + assert config.connection_config["statement_cache_size"] == 0 + await config._create_pool() # pyright: ignore[reportPrivateUsage] + assert captured_connect_kwargs["statement_cache_size"] == 0 + + alias_config = AsyncpgConfig( + connection_config={"dsn": "postgresql://localhost/test"}, driver_features={"pgbouncer": True} + ) + assert alias_config.driver_features["pooler_mode"] == "transaction" + assert alias_config.connection_config["statement_cache_size"] == 0 + + mode_alias_config = AsyncpgConfig( + connection_config={"dsn": "postgresql://localhost/test"}, driver_features={"pgbouncer_mode": True} + ) + assert mode_alias_config.driver_features["pooler_mode"] == "transaction" + assert mode_alias_config.connection_config["statement_cache_size"] == 0 + + conn = AsyncMock() + conn.is_in_transaction = lambda: True + conn.execute.return_value = "UPDATE 1" + driver = AsyncpgDriver( + cast("Any", conn), statement_config=config.statement_config, driver_features=config.driver_features + ) + + stack = StatementStack().push_execute("UPDATE items SET name = :name", {"name": "x"}) + results = await driver.execute_stack(stack) + assert results[0].result is not None + assert "prepared_statement" not in results[0].result.metadata + + +def test_asyncpg_derive_direct_config_uses_direct_dsn_and_session_mode() -> None: + """derive_direct_config should return a session-mode config targeting direct_dsn or self when unset.""" + base_config = AsyncpgConfig(connection_config={"dsn": "postgresql://localhost/test"}) + assert base_config.derive_direct_config() is base_config + + pooled_config = AsyncpgConfig( + connection_config={"dsn": "postgresql://pooler:6432/test"}, + driver_features={ + "pooler_mode": "transaction", + "direct_dsn": "postgresql://direct:5432/test", + "enable_alloydb_columnar": True, + }, + ) + direct_config = pooled_config.derive_direct_config() + assert direct_config is not pooled_config + assert direct_config.connection_config["dsn"] == "postgresql://direct:5432/test" + assert "statement_cache_size" not in direct_config.connection_config + assert direct_config.driver_features["pooler_mode"] == "session" + assert direct_config.driver_features["direct_dsn"] is None + assert direct_config.connection_config["server_settings"] == {"google_columnar_engine.enabled": "on"} + + mapping_config = AsyncpgConfig( + connection_config={"host": "pooler.internal", "port": 6432, "database": "app"}, + driver_features={"pooler_mode": "transaction", "direct_dsn": {"host": "direct.internal", "port": 5432}}, + ) + derived_mapping = mapping_config.derive_direct_config() + assert derived_mapping.connection_config["host"] == "direct.internal" + assert derived_mapping.connection_config["port"] == 5432 + assert derived_mapping.connection_config["database"] == "app" + assert derived_mapping.driver_features["pooler_mode"] == "session" + + +async def test_asyncpg_columnar_engine_session_vs_transaction() -> None: + """enable_alloydb_columnar should use server_settings in session mode and SET LOCAL in transaction mode.""" + session_config = AsyncpgConfig( + connection_config={"dsn": "postgresql://localhost/test", "server_settings": {"application_name": "sqlspec"}}, + driver_features={"pooler_mode": "session", "enable_alloydb_columnar": True}, + ) + assert session_config.connection_config["server_settings"] == { + "application_name": "sqlspec", + "google_columnar_engine.enabled": "on", + } + + tx_config = AsyncpgConfig( + connection_config={"dsn": "postgresql://localhost/test"}, + driver_features={"pooler_mode": "transaction", "enable_alloydb_columnar": True}, + ) + assert "server_settings" not in tx_config.connection_config + + tx_handle = AsyncMock() + conn = AsyncMock() + conn.is_in_transaction = lambda: False + conn.transaction = lambda: tx_handle + driver = AsyncpgDriver( + cast("Any", conn), statement_config=tx_config.statement_config, driver_features=tx_config.driver_features + ) + await driver.begin() + tx_handle.start.assert_awaited_once() + conn.execute.assert_awaited_once_with("SET LOCAL google_columnar_engine.enabled = 'on'") diff --git a/tests/unit/adapters/test_psqlpy/test_config.py b/tests/unit/adapters/test_psqlpy/test_config.py index e3bbfd3e9..79cf21092 100644 --- a/tests/unit/adapters/test_psqlpy/test_config.py +++ b/tests/unit/adapters/test_psqlpy/test_config.py @@ -13,7 +13,8 @@ build_statement_config, resolve_postgres_extension_state, ) -from sqlspec.core import StatementConfig +from sqlspec.adapters.psqlpy.driver import PsqlpyDriver +from sqlspec.core import SQL, StatementConfig class _ExtensionQueryResult: @@ -280,3 +281,75 @@ async def test_psqlpy_session_context_rolls_back_open_transaction_on_exit() -> N conn.execute.assert_any_await("ROLLBACK") assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.anyio +async def test_psqlpy_pooler_mode_transaction_disables_prepared_statements() -> None: + """Transaction pooler mode should disable statement caching and pass prepared=False on execution.""" + config = PsqlpyConfig( + connection_config={"dsn": "postgresql://localhost:6432/test"}, driver_features={"pooler_mode": "transaction"} + ) + assert config.connection_config["max_cached_statement_lifetime"] == 0 + + alias_config = PsqlpyConfig( + connection_config={"dsn": "postgresql://localhost:6432/test"}, driver_features={"pgbouncer": True} + ) + assert alias_config.driver_features["pooler_mode"] == "transaction" + assert alias_config.connection_config["max_cached_statement_lifetime"] == 0 + + conn = AsyncMock() + conn.fetch.return_value = _ExtensionQueryResult(set()) + conn.execute.return_value = "CREATE TABLE" + driver = PsqlpyDriver( + cast("PsqlpyConnection", conn), statement_config=config.statement_config, driver_features=config.driver_features + ) + + select_stmt = SQL("SELECT 1", statement_config=config.statement_config) + await driver.dispatch_execute(cast("PsqlpyConnection", conn), select_stmt) + assert conn.fetch.await_args.kwargs.get("prepared") is False + + ddl_stmt = SQL("CREATE TABLE items (id INT)", statement_config=config.statement_config) + await driver.dispatch_execute(cast("PsqlpyConnection", conn), ddl_stmt) + assert conn.execute.await_args.kwargs.get("prepared") is False + + many_stmt = SQL( + "INSERT INTO items (id) VALUES ($1)", [[1], [2]], statement_config=config.statement_config, is_many=True + ) + await driver.dispatch_execute_many(cast("PsqlpyConnection", conn), many_stmt) + assert conn.execute_many.await_args.kwargs.get("prepared") is False + + +@pytest.mark.anyio +async def test_psqlpy_derive_direct_config_and_columnar_engine() -> None: + """derive_direct_config and enable_alloydb_columnar should configure session and transaction modes.""" + base_config = PsqlpyConfig(connection_config={"dsn": "postgresql://localhost:5432/test"}) + assert base_config.derive_direct_config() is base_config + + pooled_config = PsqlpyConfig( + connection_config={"dsn": "postgresql://pooler:6432/test"}, + driver_features={ + "pooler_mode": "transaction", + "direct_dsn": "postgresql://direct:5432/test", + "enable_alloydb_columnar": True, + }, + ) + assert "options" not in pooled_config.connection_config + + direct_config = pooled_config.derive_direct_config() + assert direct_config is not pooled_config + assert direct_config.connection_config["dsn"] == "postgresql://direct:5432/test" + assert "max_cached_statement_lifetime" not in direct_config.connection_config + assert direct_config.driver_features["pooler_mode"] == "session" + assert direct_config.driver_features["direct_dsn"] is None + assert direct_config.connection_config["options"] == "-c google_columnar_engine.enabled=on" + + conn = AsyncMock() + conn.in_transaction = False + driver = PsqlpyDriver( + cast("PsqlpyConnection", conn), + statement_config=pooled_config.statement_config, + driver_features=pooled_config.driver_features, + ) + await driver.begin() + conn.execute.assert_any_await("BEGIN") + conn.execute.assert_any_await("SET LOCAL google_columnar_engine.enabled = 'on'") diff --git a/tests/unit/adapters/test_psycopg/test_config.py b/tests/unit/adapters/test_psycopg/test_config.py index 971ce872b..28a4c6e33 100644 --- a/tests/unit/adapters/test_psycopg/test_config.py +++ b/tests/unit/adapters/test_psycopg/test_config.py @@ -1,5 +1,6 @@ """Psycopg configuration tests covering statement config builders.""" +from types import SimpleNamespace from typing import Any, cast from unittest.mock import AsyncMock, MagicMock @@ -22,6 +23,7 @@ default_statement_config, resolve_postgres_extension_state, ) +from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver, PsycopgSyncDriver from sqlspec.core import SQL, StatementConfig from sqlspec.exceptions import ImproperConfigurationError @@ -685,3 +687,116 @@ async def record_hook(_: object) -> None: assert config.pg_textsearch_available is True assert config.statement_config.dialect == "pg_textsearch" assert events == ["hook", "rollback"] + + +def test_psycopg_pooler_mode_transaction_sets_prepare_threshold_none_and_alloydb_class( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """pooler_mode in {'transaction', 'statement'} sets prepare_threshold=None on sync/async configs and AlloyDB connection class.""" + sync_config = PsycopgSyncConfig( + connection_config={"conninfo": "postgresql://localhost/db"}, driver_features={"pooler_mode": "transaction"} + ) + assert sync_config.connection_config["prepare_threshold"] is None + assert sync_config.driver_features["pooler_mode"] == "transaction" + + async_config = PsycopgAsyncConfig( + connection_config={"conninfo": "postgresql://localhost/db"}, driver_features={"pooler_mode": "statement"} + ) + assert async_config.connection_config["prepare_threshold"] is None + assert async_config.driver_features["pooler_mode"] == "statement" + + alias_config = PsycopgSyncConfig( + connection_config={"conninfo": "postgresql://localhost/db"}, driver_features={"pgbouncer": True} + ) + assert alias_config.connection_config["prepare_threshold"] is None + assert alias_config.driver_features["pooler_mode"] == "transaction" + + mock_connector = MagicMock() + expected_conn = MagicMock() + expected_conn.prepare_threshold = 5 + mock_connector.connect.return_value = expected_conn + monkeypatch.setattr(psycopg_config, "ALLOYDB_CONNECTOR_INSTALLED", True) + monkeypatch.setattr(psycopg_config._psycopg_typing, "PsycopgAlloydbConnector", lambda: mock_connector) + _CapturedSyncPool.calls.clear() + monkeypatch.setattr(psycopg_config, "ConnectionPool", _CapturedSyncPool) + + alloydb_config = PsycopgSyncConfig( + connection_config={"user": "u", "password": "p", "dbname": "app"}, + driver_features={ + "enable_alloydb": True, + "alloydb_instance_uri": "projects/p/locations/r/clusters/c/instances/i", + "pooler_mode": "transaction", + }, + ) + alloydb_config._create_pool() # pyright: ignore[reportPrivateUsage] + + _conninfo, pool_kwargs = _CapturedSyncPool.calls[-1] + connection_class = pool_kwargs["connection_class"] + assert connection_class.default_prepare_threshold is None + conn = connection_class.connect("ignored") + assert conn.prepare_threshold is None + + +@pytest.mark.anyio +async def test_psycopg_enable_alloydb_columnar_and_derive_direct_config() -> None: + """enable_alloydb_columnar configures GUC by pooler_mode and derive_direct_config builds a session config.""" + session_sync = PsycopgSyncConfig( + connection_config={"conninfo": "postgresql://localhost/db", "options": "-c statement_timeout=5000"}, + driver_features={"enable_alloydb_columnar": True, "pooler_mode": "session"}, + ) + assert "-c statement_timeout=5000" in session_sync.connection_config["options"] + assert "-c google_columnar_engine.enabled=on" in session_sync.connection_config["options"] + assert session_sync.derive_direct_config() is session_sync + + txn_sync = PsycopgSyncConfig( + connection_config={"conninfo": "postgresql://pooler:6432/db"}, + driver_features={ + "enable_alloydb_columnar": True, + "pooler_mode": "transaction", + "direct_dsn": "postgresql://direct:5432/db", + }, + ) + assert "options" not in txn_sync.connection_config + assert txn_sync.connection_config["prepare_threshold"] is None + + derived_sync = txn_sync.derive_direct_config() + assert derived_sync is not txn_sync + assert derived_sync.connection_config["conninfo"] == "postgresql://direct:5432/db" + assert derived_sync.driver_features["pooler_mode"] == "session" + assert derived_sync.driver_features["direct_dsn"] is None + assert "prepare_threshold" not in derived_sync.connection_config + assert "-c google_columnar_engine.enabled=on" in derived_sync.connection_config["options"] + + txn_async = PsycopgAsyncConfig( + connection_config={"conninfo": "postgresql://pooler:6432/db"}, + driver_features={ + "enable_alloydb_columnar": True, + "pooler_mode": "transaction", + "direct_dsn": "postgresql://direct:5432/db", + }, + ) + derived_async = txn_async.derive_direct_config() + assert derived_async is not txn_async + assert derived_async.connection_config["conninfo"] == "postgresql://direct:5432/db" + assert derived_async.driver_features["pooler_mode"] == "session" + assert derived_async.driver_features["direct_dsn"] is None + assert "prepare_threshold" not in derived_async.connection_config + assert "-c google_columnar_engine.enabled=on" in derived_async.connection_config["options"] + + sync_conn = MagicMock() + sync_conn.autocommit = True + sync_conn.info = SimpleNamespace(transaction_status=0) + sync_driver = PsycopgSyncDriver( + connection=sync_conn, statement_config=txn_sync.statement_config, driver_features=txn_sync.driver_features + ) + sync_driver.begin() + sync_conn.execute.assert_called_once_with("SET LOCAL google_columnar_engine.enabled = 'on'") + + async_conn = AsyncMock() + async_conn.autocommit = True + async_conn.info = SimpleNamespace(transaction_status=0) + async_driver = PsycopgAsyncDriver( + connection=async_conn, statement_config=txn_async.statement_config, driver_features=txn_async.driver_features + ) + await async_driver.begin() + async_conn.execute.assert_awaited_once_with("SET LOCAL google_columnar_engine.enabled = 'on'") diff --git a/tests/unit/extensions/test_events/test_listener_hubs.py b/tests/unit/extensions/test_events/test_listener_hubs.py index d0f653d77..e1cf2d486 100644 --- a/tests/unit/extensions/test_events/test_listener_hubs.py +++ b/tests/unit/extensions/test_events/test_listener_hubs.py @@ -23,6 +23,7 @@ PsycopgSyncEventsBackend, PsycopgSyncHybridEventsBackend, ) +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.events import EventMessage @@ -1017,3 +1018,96 @@ async def _assert_dropped_marker_recovers(backend: Any, config: _HybridConfig) - assert "events.marker.miss" in metric_names assert "events.poll.fallback" in metric_names assert "events.dequeue.latency_ms" in metric_names + + +async def test_listener_hub_refuses_transaction_pooler_without_direct_dsn() -> None: + """All four PostgreSQL ListenerHub classes raise ImproperConfigurationError in transaction mode without direct_dsn.""" + asyncpg_cfg = _AsyncpgConfig(_AsyncpgConnection()) + asyncpg_cfg.driver_features = {"pooler_mode": "transaction", "direct_dsn": None} # type: ignore[attr-defined] + asyncpg_hub = AsyncpgListenerHub(asyncpg_cfg) # type: ignore[arg-type] + with pytest.raises(ImproperConfigurationError, match="direct_dsn"): + await asyncpg_hub.subscribe("alerts") + + psycopg_async_cfg = _PsycopgAsyncConfig(_PsycopgAsyncConnection()) + psycopg_async_cfg.driver_features = {"pooler_mode": "statement", "direct_dsn": None} # type: ignore[attr-defined] + psycopg_async_hub = PsycopgAsyncListenerHub(psycopg_async_cfg) # type: ignore[arg-type] + with pytest.raises(ImproperConfigurationError, match="direct_dsn"): + await psycopg_async_hub.subscribe("alerts") + + psycopg_sync_cfg = _PsycopgSyncConfig(_PsycopgSyncConnection()) + psycopg_sync_cfg.driver_features = {"pgbouncer": True, "direct_dsn": None} # type: ignore[attr-defined] + psycopg_sync_hub = PsycopgSyncListenerHub(psycopg_sync_cfg) # type: ignore[arg-type] + with pytest.raises(ImproperConfigurationError, match="direct_dsn"): + psycopg_sync_hub.subscribe("alerts") + + psqlpy_cfg = _PsqlpyConfig() + psqlpy_cfg.driver_features = {"pooler_mode": "transaction", "direct_dsn": None} # type: ignore[attr-defined] + psqlpy_hub = PsqlpyListenerHub(psqlpy_cfg) # type: ignore[arg-type] + with pytest.raises(ImproperConfigurationError, match="direct_dsn"): + await psqlpy_hub.subscribe("alerts") + + +async def test_listener_hub_routes_over_direct_dsn() -> None: + """All four PostgreSQL ListenerHub classes route listener connections through derive_direct_config().""" + pooled_asyncpg_conn = _AsyncpgConnection() + direct_asyncpg_conn = _AsyncpgConnection() + pooled_asyncpg_cfg = _AsyncpgConfig(pooled_asyncpg_conn) + direct_asyncpg_cfg = _AsyncpgConfig(direct_asyncpg_conn) + pooled_asyncpg_cfg.driver_features = { # type: ignore[attr-defined] + "pooler_mode": "transaction", + "direct_dsn": "postgresql://direct:5432/db", + } + pooled_asyncpg_cfg.derive_direct_config = lambda: direct_asyncpg_cfg # type: ignore[attr-defined] + + asyncpg_hub = AsyncpgListenerHub(pooled_asyncpg_cfg) # type: ignore[arg-type] + await asyncpg_hub.subscribe("alerts") + assert pooled_asyncpg_cfg.connection_enters == 0 + assert direct_asyncpg_cfg.connection_enters == 1 + assert direct_asyncpg_conn.executed == ["LISTEN alerts"] + await asyncpg_hub.shutdown() + + pooled_psycopg_async_conn = _PsycopgAsyncConnection() + direct_psycopg_async_conn = _PsycopgAsyncConnection() + pooled_psycopg_async_cfg = _PsycopgAsyncConfig(pooled_psycopg_async_conn) + direct_psycopg_async_cfg = _PsycopgAsyncConfig(direct_psycopg_async_conn) + pooled_psycopg_async_cfg.driver_features = { # type: ignore[attr-defined] + "pooler_mode": "transaction", + "direct_dsn": "postgresql://direct:5432/db", + } + pooled_psycopg_async_cfg.derive_direct_config = lambda: direct_psycopg_async_cfg # type: ignore[attr-defined] + + psycopg_async_hub = PsycopgAsyncListenerHub(pooled_psycopg_async_cfg) # type: ignore[arg-type] + await psycopg_async_hub.subscribe("alerts") + assert pooled_psycopg_async_cfg.connection_enters == 0 + assert direct_psycopg_async_cfg.connection_enters == 1 + await psycopg_async_hub.shutdown() + + pooled_psycopg_sync_conn = _PsycopgSyncConnection() + direct_psycopg_sync_conn = _PsycopgSyncConnection() + pooled_psycopg_sync_cfg = _PsycopgSyncConfig(pooled_psycopg_sync_conn) + direct_psycopg_sync_cfg = _PsycopgSyncConfig(direct_psycopg_sync_conn) + pooled_psycopg_sync_cfg.driver_features = { # type: ignore[attr-defined] + "pooler_mode": "transaction", + "direct_dsn": "postgresql://direct:5432/db", + } + pooled_psycopg_sync_cfg.derive_direct_config = lambda: direct_psycopg_sync_cfg # type: ignore[attr-defined] + + psycopg_sync_hub = PsycopgSyncListenerHub(pooled_psycopg_sync_cfg) # type: ignore[arg-type] + psycopg_sync_hub.subscribe("alerts") + assert pooled_psycopg_sync_cfg.connection_enters == 0 + assert direct_psycopg_sync_cfg.connection_enters == 1 + psycopg_sync_hub.shutdown() + + pooled_psqlpy_cfg = _PsqlpyConfig() + direct_psqlpy_cfg = _PsqlpyConfig() + pooled_psqlpy_cfg.driver_features = { # type: ignore[attr-defined] + "pooler_mode": "transaction", + "direct_dsn": "postgresql://direct:5432/db", + } + pooled_psqlpy_cfg.derive_direct_config = lambda: direct_psqlpy_cfg # type: ignore[attr-defined] + + psqlpy_hub = PsqlpyListenerHub(pooled_psqlpy_cfg) # type: ignore[arg-type] + await psqlpy_hub.subscribe("alerts") + assert pooled_psqlpy_cfg.pool.listener_calls == 0 + assert direct_psqlpy_cfg.pool.listener_calls == 1 + await psqlpy_hub.shutdown() diff --git a/tests/unit/migrations/test_migration_commands.py b/tests/unit/migrations/test_migration_commands.py index 4ca2720c4..ca10ef139 100644 --- a/tests/unit/migrations/test_migration_commands.py +++ b/tests/unit/migrations/test_migration_commands.py @@ -15,6 +15,9 @@ import pytest from sqlspec.adapters.aiosqlite.config import AiosqliteConfig +from sqlspec.adapters.asyncpg.config import AsyncpgConfig +from sqlspec.adapters.psqlpy.config import PsqlpyConfig +from sqlspec.adapters.psycopg.config import PsycopgAsyncConfig, PsycopgSyncConfig from sqlspec.adapters.sqlite.config import SqliteConfig from sqlspec.exceptions import MigrationError from sqlspec.migrations.base import BaseMigrationCommands @@ -514,3 +517,61 @@ def test_commands_contract_base_revision_file_type_default_matches_concrete_comm assert base_param.default is None assert sync_param.default is None assert async_param.default is None + + +async def test_migration_commands_route_over_direct_dsn() -> None: + """SyncMigrationCommands and AsyncMigrationCommands derive a session-mode config when direct_dsn is configured.""" + pooled_sync = PsycopgSyncConfig( + connection_config={"conninfo": "postgresql://pooler:6432/db"}, + driver_features={"pooler_mode": "transaction", "direct_dsn": "postgresql://direct:5432/db"}, + ) + sync_commands = SyncMigrationCommands(pooled_sync) + assert sync_commands.config is not pooled_sync + assert sync_commands.config.connection_config["conninfo"] == "postgresql://direct:5432/db" + assert sync_commands.config.driver_features["pooler_mode"] == "session" + assert sync_commands.config.driver_features["direct_dsn"] is None + + mock_sync_driver = MagicMock() + with ( + patch.object(sync_commands.config, "provide_session") as mock_direct_session, + patch.object(pooled_sync, "provide_session") as mock_pooled_session, + patch("sqlspec.migrations.commands.console"), + ): + mock_direct_session.return_value.__enter__.return_value = mock_sync_driver + sync_commands.tracker = MagicMock() + sync_commands.tracker.get_current_version.return_value = "0001" + assert sync_commands.current() == "0001" + mock_direct_session.assert_called_once() + mock_pooled_session.assert_not_called() + + for async_cfg in ( + AsyncpgConfig( + connection_config={"dsn": "postgresql://pooler:6432/db"}, + driver_features={"pooler_mode": "transaction", "direct_dsn": "postgresql://direct:5432/db"}, + ), + PsycopgAsyncConfig( + connection_config={"conninfo": "postgresql://pooler:6432/db"}, + driver_features={"pooler_mode": "transaction", "direct_dsn": "postgresql://direct:5432/db"}, + ), + PsqlpyConfig( + connection_config={"dsn": "postgresql://pooler:6432/db"}, + driver_features={"pooler_mode": "transaction", "direct_dsn": "postgresql://direct:5432/db"}, + ), + ): + async_commands = AsyncMigrationCommands(async_cfg) + assert async_commands.config is not async_cfg + assert async_commands.config.driver_features["pooler_mode"] == "session" + assert async_commands.config.driver_features["direct_dsn"] is None + mock_async_driver = AsyncMock() + with ( + patch.object(async_commands.config, "provide_session") as mock_direct_async_session, + patch.object(async_cfg, "provide_session") as mock_pooled_async_session, + patch("sqlspec.migrations.commands.console"), + ): + mock_direct_async_session.return_value.__aenter__.return_value = mock_async_driver + async_commands.tracker = MagicMock() + async_commands.tracker.ensure_tracking_table = AsyncMock() + async_commands.tracker.get_current_version = AsyncMock(return_value="0001") + assert await async_commands.current() == "0001" + mock_direct_async_session.assert_called_once() + mock_pooled_async_session.assert_not_called()