feat(postgres): add pooler_mode, direct_dsn session split, and AlloyDB columnar GUC across asyncpg, psycopg, and psqlpy

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

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

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())

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):

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.
 

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