sqlspec 0.38.2__tar.gz → 0.38.4__tar.gz

{sqlspec-0.38.2 → sqlspec-0.38.4}/.gitignore

@@ -61,3 +61,5 @@ specs/
 .playwright-mcp
 .geminiignore
 uv.toml
+.beads/
+.agent/

{sqlspec-0.38.2 → sqlspec-0.38.4}/.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
       - id: mixed-line-ending
       - id: trailing-whitespace
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.14.13"
+    rev: "v0.14.14"
     hooks:
       - id: ruff
         args: ["--fix"]

{sqlspec-0.38.2 → sqlspec-0.38.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlspec
-Version: 0.38.2
+Version: 0.38.4
 Summary: SQL Experiments in Python
 Project-URL: Discord, https://discord.gg/litestar
 Project-URL: Issue, https://github.com/litestar-org/sqlspec/issues/

{sqlspec-0.38.2 → sqlspec-0.38.4}/pyproject.toml

@@ -24,7 +24,7 @@ maintainers = [{ name = "Litestar Developers", email = "hello@litestar.dev" }]
 name = "sqlspec"
 readme = "README.md"
 requires-python = ">=3.10, <4.0"
-version = "0.38.2"
+version = "0.38.4"
 
 [project.urls]
 Discord = "https://discord.gg/litestar"
@@ -240,7 +240,7 @@ opt_level = "3"   # Maximum optimization (0-3)
 allow_dirty = true
 commit = false
 commit_args = "--no-verify"
-current_version = "0.38.2"
+current_version = "0.38.4"
 ignore_missing_files = false
 ignore_missing_version = false
 message = "chore(release): bump to v{new_version}"

{sqlspec-0.38.2 → sqlspec-0.38.4}/sqlspec/_typing.py

@@ -7,7 +7,7 @@ from dataclasses import dataclass
 from enum import Enum
 from typing import Any, ClassVar, Final, Literal, Protocol, cast, runtime_checkable
 
-from typing_extensions import TypeVar, dataclass_transform
+from typing_extensions import Self, TypeVar, dataclass_transform
 
 from sqlspec.utils.module_loader import dependency_flag, module_available
 
@@ -502,8 +502,8 @@ except ImportError:
         def end(self, end_time: "int | None" = None) -> None:
             return None
 
-        def __enter__(self) -> "Span":
-            return self  # type: ignore[return-value]
+        def __enter__(self) -> Self:
+            return self
 
         def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
             return None

{sqlspec-0.38.2 → sqlspec-0.38.4}/sqlspec/adapters/adbc/driver.py

@@ -7,6 +7,8 @@ database dialects, parameter style conversion, and transaction management.
 import contextlib
 from typing import TYPE_CHECKING, Any, Literal, cast
 
+from typing_extensions import Self
+
 from sqlspec.adapters.adbc._typing import AdbcSessionContext
 from sqlspec.adapters.adbc.core import (
     collect_rows,
@@ -83,7 +85,7 @@ class AdbcExceptionHandler:
     def __init__(self) -> None:
         self.pending_exception: Exception | None = None
 
-    def __enter__(self) -> "AdbcExceptionHandler":
+    def __enter__(self) -> Self:
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:

{sqlspec-0.38.2 → sqlspec-0.38.4}/sqlspec/adapters/aiosqlite/config.py

@@ -24,7 +24,7 @@ from sqlspec.utils.config_tools import normalize_connection_config
 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 @@ class AiosqliteConfig(AsyncDatabaseConfig["AiosqliteConnection", AiosqliteConnec
         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 @@ class AiosqliteConfig(AsyncDatabaseConfig["AiosqliteConnection", AiosqliteConnec
             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-0.38.2 → sqlspec-0.38.4}/sqlspec/adapters/aiosqlite/driver.py

@@ -30,6 +30,8 @@ if TYPE_CHECKING:
     from sqlspec.driver import ExecutionResult
     from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
 
+from typing_extensions import Self
+
 from sqlspec.adapters.aiosqlite._typing import AiosqliteSessionContext
 
 __all__ = ("AiosqliteCursor", "AiosqliteDriver", "AiosqliteExceptionHandler", "AiosqliteSessionContext")
@@ -81,7 +83,7 @@ class AiosqliteExceptionHandler:
     def __init__(self) -> None:
         self.pending_exception: Exception | None = None
 
-    async def __aenter__(self) -> "AiosqliteExceptionHandler":
+    async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:

{sqlspec-0.38.2 → sqlspec-0.38.4}/sqlspec/adapters/aiosqlite/pool.py

@@ -13,6 +13,7 @@ from sqlspec.utils.logging import POOL_LOGGER_NAME, get_logger, log_with_context
 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 @@ class AiosqliteConnectionPool:
         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 @@ class AiosqliteConnectionPool:
             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 @@ class AiosqliteConnectionPool:
         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 @@ class AiosqliteConnectionPool:
             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-0.38.2 → sqlspec-0.38.4}/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.extensions.events import EventRuntimeHints
 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 @@ class _AsyncmySessionFactory:
             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 @@ class AsyncmyConnectionContext:
             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 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
         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 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
         """
         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 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
         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-0.38.2 → sqlspec-0.38.4}/sqlspec/adapters/asyncmy/driver.py

@@ -39,6 +39,8 @@ if TYPE_CHECKING:
     from sqlspec.driver import ExecutionResult
     from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
 
+from typing_extensions import Self
+
 from sqlspec.adapters.asyncmy._typing import AsyncmySessionContext
 
 __all__ = ("AsyncmyCursor", "AsyncmyDriver", "AsyncmyExceptionHandler", "AsyncmySessionContext")
@@ -88,7 +90,7 @@ class AsyncmyExceptionHandler:
     def __init__(self) -> None:
         self.pending_exception: Exception | None = None
 
-    async def __aenter__(self) -> "AsyncmyExceptionHandler":
+    async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:

{sqlspec-0.38.2 → sqlspec-0.38.4}/sqlspec/adapters/asyncpg/config.py

@@ -307,12 +307,18 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async
         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 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async
         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 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async
             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-0.38.2 → sqlspec-0.38.4}/sqlspec/adapters/asyncpg/driver.py

@@ -40,6 +40,8 @@ if TYPE_CHECKING:
     from sqlspec.driver import ExecutionResult
     from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
 
+from typing_extensions import Self
+
 from sqlspec.adapters.asyncpg._typing import AsyncpgSessionContext
 
 __all__ = ("AsyncpgCursor", "AsyncpgDriver", "AsyncpgExceptionHandler", "AsyncpgSessionContext")
@@ -77,7 +79,7 @@ class AsyncpgExceptionHandler:
     def __init__(self) -> None:
         self.pending_exception: Exception | None = None
 
-    async def __aenter__(self) -> "AsyncpgExceptionHandler":
+    async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:

{sqlspec-0.38.2 → sqlspec-0.38.4}/sqlspec/adapters/bigquery/driver.py

@@ -106,7 +106,7 @@ class BigQueryExceptionHandler:
     def __init__(self) -> None:
         self.pending_exception: Exception | None = None
 
-    def __enter__(self) -> "BigQueryExceptionHandler":
+    def __enter__(self) -> "Self":
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:

{sqlspec-0.38.2 → sqlspec-0.38.4}/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 @@ class CockroachAsyncpgConfig(
         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 @@ class CockroachAsyncpgConfig(
 
     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

{sqlspec-0.38.2 → sqlspec-0.38.4}/sqlspec/adapters/cockroach_asyncpg/driver.py

@@ -5,6 +5,7 @@ import contextlib
 from typing import TYPE_CHECKING, Any, cast
 
 import asyncpg
+from typing_extensions import Self
 
 from sqlspec.adapters.asyncpg.core import create_mapped_exception, driver_profile
 from sqlspec.adapters.asyncpg.driver import AsyncpgDriver
@@ -40,7 +41,7 @@ class CockroachAsyncpgExceptionHandler:
     def __init__(self) -> None:
         self.pending_exception: Exception | None = None
 
-    async def __aenter__(self) -> "CockroachAsyncpgExceptionHandler":
+    async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:

{sqlspec-0.38.2 → sqlspec-0.38.4}/sqlspec/adapters/cockroach_psycopg/config.py

@@ -30,7 +30,7 @@ from sqlspec.extensions.events import EventRuntimeHints
 from sqlspec.utils.config_tools import normalize_connection_config
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Awaitable, Callable
 
     from sqlspec.core import StatementConfig
     from sqlspec.observability import ObservabilityConfig
@@ -82,7 +82,13 @@ class CockroachPsycopgPoolConfig(CockroachPsycopgConnectionConfig):
 
 
 class CockroachPsycopgDriverFeatures(TypedDict):
-    """CockroachDB driver feature configuration."""
+    """CockroachDB driver feature configuration.
+
+    on_connection_create: Callback executed when a connection is acquired from pool.
+        For sync: Callable[[CockroachSyncConnection], None]
+        For async: Callable[[CockroachAsyncConnection], Awaitable[None]]
+        Called after internal setup (row_factory configuration).
+    """
 
     enable_auto_retry: NotRequired[bool]
     max_retries: NotRequired[int]
@@ -94,6 +100,7 @@ class CockroachPsycopgDriverFeatures(TypedDict):
     prefer_uuid_keys: NotRequired[bool]
     json_serializer: NotRequired["Callable[[Any], str]"]
     json_deserializer: NotRequired["Callable[[str], Any]"]
+    on_connection_create: "NotRequired[Callable[..., Any]]"
     enable_events: NotRequired[bool]
     events_backend: NotRequired[str]
 
@@ -182,12 +189,18 @@ class CockroachPsycopgSyncConfig(
         driver_features.setdefault("enable_auto_retry", True)
         _ = CockroachPsycopgRetryConfig.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[[CockroachSyncConnection], 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,
@@ -228,6 +241,10 @@ class CockroachPsycopgSyncConfig(
         if autocommit_setting is not None:
             conn.autocommit = autocommit_setting
 
+        # Call user-provided callback after internal setup
+        if self._user_connection_hook is not None:
+            self._user_connection_hook(conn)
+
     def _close_pool(self) -> None:
         if not self.connection_instance:
             return
@@ -369,12 +386,18 @@ class CockroachPsycopgAsyncConfig(
         driver_features.setdefault("enable_auto_retry", True)
         _ = CockroachPsycopgRetryConfig.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[[CockroachAsyncConnection], 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,
@@ -420,6 +443,10 @@ class CockroachPsycopgAsyncConfig(
         if autocommit_setting is not None:
             await conn.set_autocommit(autocommit_setting)
 
+        # Call user-provided callback after internal setup
+        if self._user_connection_hook is not None:
+            await self._user_connection_hook(conn)
+
     async def _close_pool(self) -> None:
         if not self.connection_instance:
             return

{sqlspec-0.38.2 → sqlspec-0.38.4}/sqlspec/adapters/cockroach_psycopg/driver.py

@@ -6,6 +6,7 @@ import time
 from typing import TYPE_CHECKING, Any, cast
 
 import psycopg
+from typing_extensions import Self
 
 from sqlspec.adapters.cockroach_psycopg._typing import (
     CockroachAsyncConnection,
@@ -57,7 +58,7 @@ class CockroachPsycopgSyncExceptionHandler:
     def __init__(self) -> None:
         self.pending_exception: Exception | None = None
 
-    def __enter__(self) -> "CockroachPsycopgSyncExceptionHandler":
+    def __enter__(self) -> Self:
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
@@ -80,7 +81,7 @@ class CockroachPsycopgAsyncExceptionHandler:
     def __init__(self) -> None:
         self.pending_exception: Exception | None = None
 
-    async def __aenter__(self) -> "CockroachPsycopgAsyncExceptionHandler":
+    async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool: