feat: add on_connection_create hooks and fix DuckDB variable persistence (#342)

- Adds the already implemented `on_connection_create` callback support to all adapter configs, allowing custom connection initialization logic  (#340)
- Fixes DuckDB `SET VARIABLE` persistence across `execute()` calls (#341)

Author: cofin

sqlspec/adapters/aiosqlite/config.py

@@ -24,7 +24,7 @@
 from sqlspec.utils.logging import get_logger
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Awaitable, Callable
 
     from sqlspec.core import StatementConfig
     from sqlspec.observability import ObservabilityConfig
@@ -68,6 +68,9 @@ class AiosqliteDriverFeatures(TypedDict):
         Defaults to sqlspec.utils.serializers.to_json.
     json_deserializer: Custom JSON deserializer function.
         Defaults to sqlspec.utils.serializers.from_json.
+    on_connection_create: Async callback executed when a connection is created.
+        Receives the raw aiosqlite connection for low-level driver configuration.
+        Runs after internal setup (PRAGMA optimizations).
     enable_events: Enable database event channel support.
         Defaults to True when extension_config["events"] is configured.
         Provides pub/sub capabilities via table-backed queue (SQLite has no native pub/sub).
@@ -81,6 +84,7 @@ class AiosqliteDriverFeatures(TypedDict):
     enable_custom_adapters: NotRequired[bool]
     json_serializer: "NotRequired[Callable[[Any], str]]"
     json_deserializer: "NotRequired[Callable[[str], Any]]"
+    on_connection_create: "NotRequired[Callable[[AiosqliteConnection], Awaitable[None]]]"
     enable_events: NotRequired[bool]
     events_backend: NotRequired[str]
 
@@ -191,12 +195,18 @@ def __init__(
         statement_config = statement_config or default_statement_config
         statement_config, driver_features = apply_driver_features(statement_config, driver_features)
 
+        # Extract user connection hook before storing driver_features
+        features_dict = dict(driver_features) if driver_features else {}
+        self._user_connection_hook: Callable[[AiosqliteConnection], Awaitable[None]] | None = features_dict.pop(
+            "on_connection_create", None
+        )
+
         super().__init__(
             connection_config=config_dict,
             connection_instance=connection_instance,
             migration_config=migration_config,
             statement_config=statement_config,
-            driver_features=driver_features,
+            driver_features=features_dict,
             bind_key=bind_key,
             extension_config=extension_config,
             observability_config=observability_config,
@@ -258,6 +268,7 @@ async def _create_pool(self) -> AiosqliteConnectionPool:
             connect_timeout=connect_timeout,
             idle_timeout=idle_timeout,
             operation_timeout=operation_timeout,
+            on_connection_create=self._user_connection_hook,
         )
 
         if self.driver_features.get("enable_custom_adapters", False):

sqlspec/adapters/aiosqlite/pool.py

@@ -13,6 +13,7 @@
 from sqlspec.utils.uuids import uuid4
 
 if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
     from types import TracebackType
 
     from sqlspec.adapters.aiosqlite._typing import AiosqliteConnection
@@ -178,6 +179,7 @@ class AiosqliteConnectionPool:
         "_idle_timeout",
         "_lock_instance",
         "_min_size",
+        "_on_connection_create",
         "_operation_timeout",
         "_pool_id",
         "_pool_size",
@@ -195,6 +197,7 @@ def __init__(
         idle_timeout: float = 24 * 60 * 60,
         operation_timeout: float = 10.0,
         health_check_interval: float = 30.0,
+        on_connection_create: "Callable[[AiosqliteConnection], Awaitable[None]] | None" = None,
     ) -> None:
         """Initialize connection pool.
 
@@ -206,6 +209,7 @@ def __init__(
             idle_timeout: Maximum time a connection can remain idle
             operation_timeout: Maximum time for connection operations
             health_check_interval: Seconds of idle time before running health check
+            on_connection_create: Async callback executed when connection is created
         """
         self._connection_parameters = connection_parameters
         self._pool_size = pool_size
@@ -214,6 +218,7 @@ def __init__(
         self._idle_timeout = idle_timeout
         self._operation_timeout = operation_timeout
         self._health_check_interval = health_check_interval
+        self._on_connection_create = on_connection_create
 
         self._connection_registry: dict[str, AiosqlitePoolConnection] = {}
         self._wal_initialized = False
@@ -324,6 +329,10 @@ async def _create_connection(self) -> AiosqlitePoolConnection:
             await connection.execute("PRAGMA busy_timeout = 30000")
             await connection.commit()
 
+        # Call user-provided callback after internal setup
+        if self._on_connection_create is not None:
+            await self._on_connection_create(connection)
+
         pool_connection = AiosqlitePoolConnection(connection)
         pool_connection.mark_as_idle()
 

sqlspec/adapters/asyncmy/config.py

@@ -1,6 +1,7 @@
 """Asyncmy database configuration."""
 
 from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, cast
+from weakref import WeakSet
 
 import asyncmy
 from asyncmy.cursors import Cursor, DictCursor  # pyright: ignore
@@ -16,7 +17,7 @@
 from sqlspec.utils.config_tools import normalize_connection_config
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Awaitable, Callable
 
     from asyncmy.cursors import Cursor, DictCursor  # pyright: ignore
     from asyncmy.pool import Pool  # pyright: ignore
@@ -71,6 +72,9 @@ class AsyncmyDriverFeatures(TypedDict):
     json_deserializer: Custom JSON deserializer function.
         Defaults to sqlspec.utils.serializers.from_json.
         Use for performance (orjson) or custom decoding.
+    on_connection_create: Async callback executed when a connection is acquired from pool.
+        Receives the raw asyncmy connection for low-level driver configuration.
+        Called exactly once per physical connection using WeakSet tracking.
     enable_events: Enable database event channel support.
         Defaults to True when extension_config["events"] is configured.
         Provides pub/sub capabilities via table-backed queue (MySQL/MariaDB have no native pub/sub).
@@ -83,6 +87,7 @@ class AsyncmyDriverFeatures(TypedDict):
 
     json_serializer: NotRequired["Callable[[Any], str]"]
     json_deserializer: NotRequired["Callable[[str], Any]"]
+    on_connection_create: "NotRequired[Callable[[AsyncmyConnection], Awaitable[None]]]"
     enable_events: NotRequired[bool]
     events_backend: NotRequired[str]
 
@@ -101,7 +106,9 @@ async def acquire_connection(self) -> "AsyncmyConnection":
             self._config.connection_instance = pool
         ctx = pool.acquire()
         self._ctx = ctx
-        return cast("AsyncmyConnection", await ctx.__aenter__())
+        connection = cast("AsyncmyConnection", await ctx.__aenter__())
+        await self._config._ensure_connection_initialized(connection)  # pyright: ignore[reportPrivateUsage]
+        return connection
 
     async def release_connection(self, _conn: "AsyncmyConnection") -> None:
         if self._ctx is not None:
@@ -125,7 +132,9 @@ async def __aenter__(self) -> AsyncmyConnection:
             self._config.connection_instance = pool
         ctx = pool.acquire()
         self._ctx = ctx
-        return cast("AsyncmyConnection", await ctx.__aenter__())
+        connection = cast("AsyncmyConnection", await ctx.__aenter__())
+        await self._config._ensure_connection_initialized(connection)  # pyright: ignore[reportPrivateUsage]
+        return connection
 
     async def __aexit__(
         self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: Any
@@ -181,12 +190,20 @@ def __init__(
         statement_config = statement_config or default_statement_config
         statement_config, driver_features = apply_driver_features(statement_config, driver_features)
 
+        # Extract user connection hook before storing driver_features
+        features_dict = dict(driver_features) if driver_features else {}
+        self._user_connection_hook: Callable[[AsyncmyConnection], Awaitable[None]] | None = features_dict.pop(
+            "on_connection_create", None
+        )
+        # Track initialized connections to ensure callback runs exactly once per physical connection
+        self._initialized_connections: WeakSet[Any] = WeakSet()
+
         super().__init__(
             connection_config=connection_config,
             connection_instance=connection_instance,
             migration_config=migration_config,
             statement_config=statement_config,
-            driver_features=driver_features,
+            driver_features=features_dict,
             bind_key=bind_key,
             extension_config=extension_config,
             observability_config=observability_config,
@@ -205,6 +222,17 @@ async def _create_pool(self) -> "AsyncmyPool":
         """
         return cast("AsyncmyPool", await asyncmy.create_pool(**dict(self.connection_config)))
 
+    async def _ensure_connection_initialized(self, connection: "AsyncmyConnection") -> None:
+        """Ensure connection callback has been called exactly once for this connection.
+
+        Uses WeakSet tracking to ensure the callback runs once per physical connection.
+        """
+        if self._user_connection_hook is None:
+            return
+        if connection not in self._initialized_connections:
+            await self._user_connection_hook(connection)
+            self._initialized_connections.add(connection)
+
     async def _close_pool(self) -> None:
         """Close the actual async connection pool."""
         if self.connection_instance:
@@ -226,7 +254,9 @@ async def create_connection(self) -> AsyncmyConnection:
         if pool is None:
             pool = await self.create_pool()
             self.connection_instance = pool
-        return cast("AsyncmyConnection", await pool.acquire())
+        connection = cast("AsyncmyConnection", await pool.acquire())
+        await self._ensure_connection_initialized(connection)
+        return connection
 
     def provide_connection(self, *args: Any, **kwargs: Any) -> "AsyncmyConnectionContext":
         """Provide an async connection context manager.

sqlspec/adapters/asyncpg/config.py

@@ -307,12 +307,18 @@ def __init__(
         statement_config = statement_config or default_statement_config
         statement_config, driver_features = apply_driver_features(statement_config, driver_features)
 
+        # Extract user connection hook before storing driver_features
+        features_dict = dict(driver_features) if driver_features else {}
+        self._user_connection_hook: Callable[[AsyncpgConnection], Awaitable[None]] | None = features_dict.pop(
+            "on_connection_create", None
+        )
+
         super().__init__(
             connection_config=normalize_connection_config(connection_config),
             connection_instance=connection_instance,
             migration_config=migration_config,
             statement_config=statement_config,
-            driver_features=driver_features,
+            driver_features=features_dict,
             bind_key=bind_key,
             extension_config=extension_config,
             observability_config=observability_config,
@@ -432,7 +438,7 @@ async def _create_pool(self) -> "Pool[Record]":
         return await asyncpg_create_pool(**config)
 
     async def _init_connection(self, connection: "AsyncpgConnection") -> None:
-        """Initialize connection with JSON codecs and pgvector support.
+        """Initialize connection with JSON codecs, pgvector support, and user callback.
 
         Args:
             connection: AsyncPG connection to initialize.
@@ -456,6 +462,10 @@ async def _init_connection(self, connection: "AsyncpgConnection") -> None:
             if self._pgvector_available:
                 await register_pgvector_support(connection)
 
+        # Call user-provided callback after internal setup
+        if self._user_connection_hook is not None:
+            await self._user_connection_hook(connection)
+
     async def _close_pool(self) -> None:
         """Close the actual async connection pool and cleanup connectors."""
         if self.connection_instance:

sqlspec/adapters/cockroach_asyncpg/config.py

@@ -64,7 +64,12 @@ class CockroachAsyncpgPoolConfig(CockroachAsyncpgConnectionConfig):
 
 
 class CockroachAsyncpgDriverFeatures(TypedDict):
-    """Driver feature flags for CockroachDB AsyncPG adapter."""
+    """Driver feature flags for CockroachDB AsyncPG adapter.
+
+    on_connection_create: Async callback executed when a connection is acquired from pool.
+        Receives the raw asyncpg connection for low-level driver configuration.
+        Called after internal setup (JSON codecs, pgvector registration).
+    """
 
     enable_auto_retry: NotRequired[bool]
     max_retries: NotRequired[int]
@@ -77,6 +82,7 @@ class CockroachAsyncpgDriverFeatures(TypedDict):
     json_deserializer: NotRequired["Callable[[str], Any]"]
     enable_json_codecs: NotRequired[bool]
     enable_pgvector: NotRequired[bool]
+    on_connection_create: "NotRequired[Callable[[CockroachAsyncpgConnection], Awaitable[None]]]"
     enable_events: NotRequired[bool]
     events_backend: NotRequired[str]
 
@@ -163,12 +169,18 @@ def __init__(
         driver_features.setdefault("enable_auto_retry", True)
         _ = CockroachAsyncpgRetryConfig.from_features(driver_features)
 
+        # Extract user connection hook before storing driver_features
+        features_dict = dict(driver_features) if driver_features else {}
+        self._user_connection_hook: Callable[[CockroachAsyncpgConnection], Awaitable[None]] | None = features_dict.pop(
+            "on_connection_create", None
+        )
+
         super().__init__(
             connection_config=connection_config,
             connection_instance=connection_instance,
             migration_config=migration_config,
             statement_config=statement_config,
-            driver_features=driver_features,
+            driver_features=features_dict,
             bind_key=bind_key,
             extension_config=extension_config,
             observability_config=observability_config,
@@ -177,8 +189,14 @@ def __init__(
 
     async def _create_pool(self) -> "CockroachAsyncpgPool":
         config = build_connection_config(self.connection_config)
+        config.setdefault("init", self._init_connection)
         return await asyncpg_create_pool(**config)
 
+    async def _init_connection(self, connection: "CockroachAsyncpgConnection") -> None:
+        """Initialize connection with user callback if provided."""
+        if self._user_connection_hook is not None:
+            await self._user_connection_hook(connection)
+
     async def _close_pool(self) -> None:
         if not self.connection_instance:
             return