sqlspec 0.16.2__py3-none-any.whl → 0.17.0__py3-none-any.whl

sqlspec/adapters/duckdb/config.py

@@ -1,38 +1,24 @@
 """DuckDB database configuration with connection pooling."""
-# ruff: noqa: D107 W293 RUF100 S110 PLR0913 FA100 BLE001 UP037 COM812 ARG002
 
-import logging
-import threading
-import time
 from collections.abc import Sequence
-from contextlib import contextmanager, suppress
-from typing import TYPE_CHECKING, Any, Final, Optional, TypedDict, cast
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, TypedDict, Union, cast
 
-import duckdb
 from typing_extensions import NotRequired
 
 from sqlspec.adapters.duckdb._types import DuckDBConnection
 from sqlspec.adapters.duckdb.driver import DuckDBCursor, DuckDBDriver, duckdb_statement_config
+from sqlspec.adapters.duckdb.pool import DuckDBConnectionPool
 from sqlspec.config import SyncDatabaseConfig
 
 if TYPE_CHECKING:
-    from collections.abc import Generator
-    from typing import Callable, ClassVar, Union
+    from collections.abc import Callable, Generator
 
     from sqlspec.core.statement import StatementConfig
 
-
-logger = logging.getLogger(__name__)
-
-DEFAULT_MIN_POOL: Final[int] = 1
-DEFAULT_MAX_POOL: Final[int] = 4
-POOL_TIMEOUT: Final[float] = 30.0
-POOL_RECYCLE: Final[int] = 86400
-
 __all__ = (
     "DuckDBConfig",
     "DuckDBConnectionParams",
-    "DuckDBConnectionPool",
     "DuckDBDriverFeatures",
     "DuckDBExtensionConfig",
     "DuckDBPoolParams",
@@ -40,210 +26,6 @@ __all__ = (
 )
 
 
-class DuckDBConnectionPool:
-    """Thread-local connection manager for DuckDB with performance optimizations.
-
-    Uses thread-local storage to ensure each thread gets its own DuckDB connection,
-    preventing the thread-safety issues that cause segmentation faults when
-    multiple cursors share the same connection concurrently.
-
-    This design trades traditional pooling for thread safety, which is essential
-    for DuckDB since connections and cursors are not thread-safe.
-    """
-
-    __slots__ = (
-        "_connection_config",
-        "_connection_times",
-        "_created_connections",
-        "_extensions",
-        "_lock",
-        "_on_connection_create",
-        "_recycle",
-        "_secrets",
-        "_thread_local",
-    )
-
-    def __init__(  # noqa: PLR0913
-        self,
-        connection_config: "dict[str, Any]",  # noqa: UP037
-        pool_min_size: int = DEFAULT_MIN_POOL,
-        pool_max_size: int = DEFAULT_MAX_POOL,
-        pool_timeout: float = POOL_TIMEOUT,
-        pool_recycle_seconds: int = POOL_RECYCLE,
-        extensions: "Optional[list[dict[str, Any]]]" = None,  # noqa: FA100, UP037
-        secrets: "Optional[list[dict[str, Any]]]" = None,  # noqa: FA100, UP037
-        on_connection_create: "Optional[Callable[[DuckDBConnection], None]]" = None,  # noqa: FA100
-    ) -> None:
-        """Initialize the thread-local connection manager."""
-        self._connection_config = connection_config
-        self._recycle = pool_recycle_seconds
-        self._extensions = extensions or []
-        self._secrets = secrets or []
-        self._on_connection_create = on_connection_create
-        self._thread_local = threading.local()
-        self._lock = threading.RLock()
-        self._created_connections = 0
-        self._connection_times: "dict[int, float]" = {}
-
-    def _create_connection(self) -> DuckDBConnection:
-        """Create a new DuckDB connection with extensions and secrets."""
-        connect_parameters = {}
-        config_dict = {}
-
-        for key, value in self._connection_config.items():
-            if key in {"database", "read_only"}:
-                connect_parameters[key] = value
-            else:
-                config_dict[key] = value
-
-        if config_dict:
-            connect_parameters["config"] = config_dict
-
-        connection = duckdb.connect(**connect_parameters)
-
-        for ext_config in self._extensions:
-            ext_name = ext_config.get("name")
-            if not ext_name:
-                continue
-
-            install_kwargs = {}
-            if "version" in ext_config:
-                install_kwargs["version"] = ext_config["version"]
-            if "repository" in ext_config:
-                install_kwargs["repository"] = ext_config["repository"]
-            if ext_config.get("force_install", False):
-                install_kwargs["force_install"] = True
-
-            try:
-                if install_kwargs:
-                    connection.install_extension(ext_name, **install_kwargs)
-                connection.load_extension(ext_name)
-            except Exception:  # noqa: BLE001, S110
-                pass
-
-        for secret_config in self._secrets:
-            secret_type = secret_config.get("secret_type")
-            secret_name = secret_config.get("name")
-            secret_value = secret_config.get("value")
-
-            if not (secret_type and secret_name and secret_value):
-                continue
-
-            value_pairs = []
-            for key, value in secret_value.items():
-                escaped_value = str(value).replace("'", "''")
-                value_pairs.append(f"'{key}' = '{escaped_value}'")
-            value_string = ", ".join(value_pairs)
-            scope_clause = ""
-            if "scope" in secret_config:
-                scope_clause = f" SCOPE '{secret_config['scope']}'"
-
-            sql = f"""  # noqa: S608
-                CREATE SECRET {secret_name} (
-                    TYPE {secret_type},
-                    {value_string}
-                ){scope_clause}
-            """
-            with suppress(Exception):
-                connection.execute(sql)
-
-        if self._on_connection_create:
-            with suppress(Exception):
-                self._on_connection_create(connection)
-
-        conn_id = id(connection)
-        with self._lock:
-            self._created_connections += 1
-            self._connection_times[conn_id] = time.time()
-
-        return connection
-
-    def _get_thread_connection(self) -> DuckDBConnection:
-        """Get or create a connection for the current thread.
-
-        Each thread gets its own dedicated DuckDB connection to prevent
-        thread-safety issues with concurrent cursor operations.
-        """
-        if not hasattr(self._thread_local, "connection"):
-            self._thread_local.connection = self._create_connection()
-            self._thread_local.created_at = time.time()
-
-        # Check if connection needs recycling
-        if self._recycle > 0 and time.time() - self._thread_local.created_at > self._recycle:
-            with suppress(Exception):
-                self._thread_local.connection.close()
-            self._thread_local.connection = self._create_connection()
-            self._thread_local.created_at = time.time()
-
-        return cast("DuckDBConnection", self._thread_local.connection)
-
-    def _close_thread_connection(self) -> None:
-        """Close the connection for the current thread."""
-        if hasattr(self._thread_local, "connection"):
-            with suppress(Exception):
-                self._thread_local.connection.close()
-            del self._thread_local.connection
-            if hasattr(self._thread_local, "created_at"):
-                del self._thread_local.created_at
-
-    def _is_connection_alive(self, connection: DuckDBConnection) -> bool:
-        """Check if a connection is still alive and usable.
-
-        Args:
-            connection: Connection to check
-
-        Returns:
-            True if connection is alive, False otherwise
-        """
-        try:
-            cursor = connection.cursor()
-            cursor.close()
-        except Exception:
-            return False
-        return True
-
-    @contextmanager
-    def get_connection(self) -> "Generator[DuckDBConnection, None, None]":
-        """Get a thread-local connection.
-
-        Each thread gets its own dedicated DuckDB connection to prevent
-        thread-safety issues with concurrent cursor operations.
-
-        Yields:
-            DuckDBConnection: A thread-local connection.
-        """
-        connection = self._get_thread_connection()
-        try:
-            yield connection
-        except Exception:
-            # On error, close and recreate connection for this thread
-            self._close_thread_connection()
-            raise
-
-    def close(self) -> None:
-        """Close the thread-local connection if it exists."""
-        self._close_thread_connection()
-
-    def size(self) -> int:
-        """Get current pool size (always 1 for thread-local)."""
-        return 1 if hasattr(self._thread_local, "connection") else 0
-
-    def checked_out(self) -> int:
-        """Get number of checked out connections (always 0 for thread-local)."""
-        return 0
-
-    def acquire(self) -> DuckDBConnection:
-        """Acquire a thread-local connection.
-
-        Each thread gets its own dedicated DuckDB connection to prevent
-        thread-safety issues with concurrent cursor operations.
-
-        Returns:
-            DuckDBConnection: A thread-local connection
-        """
-        return self._get_thread_connection()
-
-
 class DuckDBConnectionParams(TypedDict, total=False):
     """DuckDB connection parameters."""
 
@@ -398,16 +180,9 @@ class DuckDBConfig(SyncDatabaseConfig[DuckDBConnection, DuckDBConnectionPool, Du
             and k not in {"pool_min_size", "pool_max_size", "pool_timeout", "pool_recycle_seconds", "extra"}
         }
 
-    def _get_pool_config_dict(self) -> "dict[str, Any]":
-        """Get pool configuration as plain dict for pool creation."""
-        return {
-            k: v
-            for k, v in self.pool_config.items()
-            if v is not None and k in {"pool_min_size", "pool_max_size", "pool_timeout", "pool_recycle_seconds"}
-        }
-
     def _create_pool(self) -> DuckDBConnectionPool:
-        """Create the DuckDB connection pool."""
+        """Create connection pool from configuration."""
+        connection_config = self._get_connection_config_dict()
 
         extensions = self.driver_features.get("extensions", None)
         secrets = self.driver_features.get("secrets", None)
@@ -423,12 +198,13 @@ class DuckDBConfig(SyncDatabaseConfig[DuckDBConnection, DuckDBConnectionPool, Du
                 on_connection_create(conn)
 
             pool_callback = wrapped_callback
-        conf = {"extensions": extensions_dicts, "secrets": secrets_dicts, "on_connection_create": pool_callback}
 
         return DuckDBConnectionPool(
-            connection_config=self._get_connection_config_dict(),
-            **conf,  # type: ignore[arg-type]
-            **self._get_pool_config_dict(),
+            connection_config=connection_config,
+            extensions=extensions_dicts,
+            secrets=secrets_dicts,
+            on_connection_create=pool_callback,
+            **self.pool_config,  # Pass all pool_config as kwargs to be filtered by the pool
         )
 
     def _close_pool(self) -> None:

sqlspec/adapters/duckdb/pool.py

@@ -0,0 +1,243 @@
+"""DuckDB connection pool with thread-local connections."""
+
+import logging
+import threading
+import time
+from contextlib import contextmanager, suppress
+from typing import TYPE_CHECKING, Any, Final, Optional, cast
+
+import duckdb
+
+from sqlspec.adapters.duckdb._types import DuckDBConnection
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Generator
+
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MIN_POOL: Final[int] = 1
+DEFAULT_MAX_POOL: Final[int] = 4
+POOL_TIMEOUT: Final[float] = 30.0
+POOL_RECYCLE: Final[int] = 86400
+
+__all__ = ("DuckDBConnectionPool",)
+
+
+class DuckDBConnectionPool:
+    """Thread-local connection manager for DuckDB with performance optimizations.
+
+    Uses thread-local storage to ensure each thread gets its own DuckDB connection,
+    preventing the thread-safety issues that cause segmentation faults when
+    multiple cursors share the same connection concurrently.
+
+    This design trades traditional pooling for thread safety, which is essential
+    for DuckDB since connections and cursors are not thread-safe.
+    """
+
+    __slots__ = (
+        "_connection_config",
+        "_connection_times",
+        "_created_connections",
+        "_extensions",
+        "_lock",
+        "_on_connection_create",
+        "_recycle",
+        "_secrets",
+        "_thread_local",
+    )
+
+    def __init__(
+        self,
+        connection_config: "dict[str, Any]",
+        pool_recycle_seconds: int = POOL_RECYCLE,
+        extensions: "Optional[list[dict[str, Any]]]" = None,
+        secrets: "Optional[list[dict[str, Any]]]" = None,
+        on_connection_create: "Optional[Callable[[DuckDBConnection], None]]" = None,
+        **kwargs: Any,  # Accept and ignore additional parameters for compatibility
+    ) -> None:
+        """Initialize the thread-local connection manager.
+
+        Args:
+            connection_config: DuckDB connection configuration
+            pool_recycle_seconds: Connection recycle time in seconds
+            extensions: List of extensions to install/load
+            secrets: List of secrets to create
+            on_connection_create: Callback executed when connection is created
+            **kwargs: Additional parameters ignored for compatibility
+        """
+        self._connection_config = connection_config
+        self._recycle = pool_recycle_seconds
+        self._extensions = extensions or []
+        self._secrets = secrets or []
+        self._on_connection_create = on_connection_create
+        self._thread_local = threading.local()
+        self._lock = threading.RLock()
+        self._created_connections = 0
+        self._connection_times: dict[int, float] = {}
+
+    def _create_connection(self) -> DuckDBConnection:
+        """Create a new DuckDB connection with extensions and secrets."""
+        connect_parameters = {}
+        config_dict = {}
+
+        for key, value in self._connection_config.items():
+            if key in {"database", "read_only"}:
+                connect_parameters[key] = value
+            else:
+                config_dict[key] = value
+
+        if config_dict:
+            connect_parameters["config"] = config_dict
+
+        connection = duckdb.connect(**connect_parameters)
+
+        for ext_config in self._extensions:
+            ext_name = ext_config.get("name")
+            if not ext_name:
+                continue
+
+            install_kwargs = {}
+            if "version" in ext_config:
+                install_kwargs["version"] = ext_config["version"]
+            if "repository" in ext_config:
+                install_kwargs["repository"] = ext_config["repository"]
+            if ext_config.get("force_install", False):
+                install_kwargs["force_install"] = True
+
+            try:
+                if install_kwargs:
+                    connection.install_extension(ext_name, **install_kwargs)
+                connection.load_extension(ext_name)
+            except Exception as e:
+                logger.debug("Failed to load DuckDB extension %s: %s", ext_name, e)
+
+        for secret_config in self._secrets:
+            secret_type = secret_config.get("secret_type")
+            secret_name = secret_config.get("name")
+            secret_value = secret_config.get("value")
+
+            if not (secret_type and secret_name and secret_value):
+                continue
+
+            value_pairs = []
+            for key, value in secret_value.items():
+                escaped_value = str(value).replace("'", "''")
+                value_pairs.append(f"'{key}' = '{escaped_value}'")
+            value_string = ", ".join(value_pairs)
+            scope_clause = ""
+            if "scope" in secret_config:
+                scope_clause = f" SCOPE '{secret_config['scope']}'"
+
+            sql = f"""
+                CREATE SECRET {secret_name} (
+                    TYPE {secret_type},
+                    {value_string}
+                ){scope_clause}
+            """
+            with suppress(Exception):
+                connection.execute(sql)
+
+        if self._on_connection_create:
+            with suppress(Exception):
+                self._on_connection_create(connection)
+
+        conn_id = id(connection)
+        with self._lock:
+            self._created_connections += 1
+            self._connection_times[conn_id] = time.time()
+
+        return connection
+
+    def _get_thread_connection(self) -> DuckDBConnection:
+        """Get or create a connection for the current thread.
+
+        Each thread gets its own dedicated DuckDB connection to prevent
+        thread-safety issues with concurrent cursor operations.
+        """
+        if not hasattr(self._thread_local, "connection"):
+            self._thread_local.connection = self._create_connection()
+            self._thread_local.created_at = time.time()
+
+        # Check if connection needs recycling
+        if self._recycle > 0 and time.time() - self._thread_local.created_at > self._recycle:
+            with suppress(Exception):
+                self._thread_local.connection.close()
+            self._thread_local.connection = self._create_connection()
+            self._thread_local.created_at = time.time()
+
+        return cast("DuckDBConnection", self._thread_local.connection)
+
+    def _close_thread_connection(self) -> None:
+        """Close the connection for the current thread."""
+        if hasattr(self._thread_local, "connection"):
+            with suppress(Exception):
+                self._thread_local.connection.close()
+            del self._thread_local.connection
+            if hasattr(self._thread_local, "created_at"):
+                del self._thread_local.created_at
+
+    def _is_connection_alive(self, connection: DuckDBConnection) -> bool:
+        """Check if a connection is still alive and usable.
+
+        Args:
+            connection: Connection to check
+
+        Returns:
+            True if connection is alive, False otherwise
+        """
+        try:
+            cursor = connection.cursor()
+            cursor.close()
+        except Exception:
+            return False
+        return True
+
+    @contextmanager
+    def get_connection(self) -> "Generator[DuckDBConnection, None, None]":
+        """Get a thread-local connection.
+
+        Each thread gets its own dedicated DuckDB connection to prevent
+        thread-safety issues with concurrent cursor operations.
+
+        Yields:
+            DuckDBConnection: A thread-local connection.
+        """
+        connection = self._get_thread_connection()
+        try:
+            yield connection
+        except Exception:
+            # On error, close and recreate connection for this thread
+            self._close_thread_connection()
+            raise
+
+    def close(self) -> None:
+        """Close the thread-local connection if it exists."""
+        self._close_thread_connection()
+
+    def size(self) -> int:
+        """Get current pool size (always 1 for thread-local)."""
+        return 1 if hasattr(self._thread_local, "connection") else 0
+
+    def checked_out(self) -> int:
+        """Get number of checked out connections (always 0 for thread-local)."""
+        return 0
+
+    def acquire(self) -> DuckDBConnection:
+        """Acquire a thread-local connection.
+
+        Each thread gets its own dedicated DuckDB connection to prevent
+        thread-safety issues with concurrent cursor operations.
+
+        Returns:
+            DuckDBConnection: A thread-local connection
+        """
+        return self._get_thread_connection()
+
+    def release(self, connection: DuckDBConnection) -> None:
+        """Release a connection (no-op for thread-local connections).
+
+        Args:
+            connection: The connection to release (ignored)
+        """
+        # No-op: thread-local connections are managed per-thread

sqlspec/adapters/sqlite/__init__.py

@@ -3,11 +3,13 @@
 from sqlspec.adapters.sqlite._types import SqliteConnection
 from sqlspec.adapters.sqlite.config import SqliteConfig, SqliteConnectionParams
 from sqlspec.adapters.sqlite.driver import SqliteCursor, SqliteDriver, SqliteExceptionHandler, sqlite_statement_config
+from sqlspec.adapters.sqlite.pool import SqliteConnectionPool
 
 __all__ = (
     "SqliteConfig",
     "SqliteConnection",
     "SqliteConnectionParams",
+    "SqliteConnectionPool",
     "SqliteCursor",
     "SqliteDriver",
     "SqliteExceptionHandler",

sqlspec/adapters/sqlite/config.py

@@ -1,7 +1,6 @@
 """SQLite database configuration with thread-local connections."""
 
-import sqlite3
-import threading
+import uuid
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, ClassVar, Optional, TypedDict, Union, cast
 
@@ -9,6 +8,7 @@ from typing_extensions import NotRequired
 
 from sqlspec.adapters.sqlite._types import SqliteConnection
 from sqlspec.adapters.sqlite.driver import SqliteCursor, SqliteDriver, sqlite_statement_config
+from sqlspec.adapters.sqlite.pool import SqliteConnectionPool
 from sqlspec.config import SyncDatabaseConfig
 
 if TYPE_CHECKING:
@@ -30,118 +30,7 @@ class SqliteConnectionParams(TypedDict, total=False):
     uri: NotRequired[bool]
 
 
-__all__ = ("SqliteConfig", "SqliteConnectionParams", "SqliteConnectionPool")
-
-
-class SqliteConnectionPool:
-    """Thread-local connection manager for SQLite.
-
-    SQLite connections aren't thread-safe, so we use thread-local storage
-    to ensure each thread has its own connection. This is simpler and more
-    efficient than a traditional pool for SQLite's constraints.
-    """
-
-    __slots__ = ("_connection_parameters", "_enable_optimizations", "_thread_local")
-
-    def __init__(
-        self,
-        connection_parameters: "dict[str, Any]",
-        enable_optimizations: bool = True,
-        **kwargs: Any,  # Accept and ignore pool parameters for compatibility
-    ) -> None:
-        """Initialize the thread-local connection manager.
-
-        Args:
-            connection_parameters: SQLite connection parameters
-            enable_optimizations: Whether to apply performance PRAGMAs
-            **kwargs: Ignored pool parameters for compatibility
-        """
-        self._connection_parameters = connection_parameters
-        self._thread_local = threading.local()
-        self._enable_optimizations = enable_optimizations
-
-    def _create_connection(self) -> SqliteConnection:
-        """Create a new SQLite connection with optimizations."""
-        connection = sqlite3.connect(**self._connection_parameters)
-
-        # Only apply optimizations if requested and not in-memory
-        if self._enable_optimizations:
-            database = self._connection_parameters.get("database", ":memory:")
-            is_memory = database == ":memory:" or database.startswith("file::memory:")
-
-            if not is_memory:
-                # WAL mode doesn't work with in-memory databases
-                connection.execute("PRAGMA journal_mode = WAL")
-                # Set busy timeout for better concurrent access
-                connection.execute("PRAGMA busy_timeout = 5000")
-                connection.execute("PRAGMA optimize")
-            # These work for all database types
-            connection.execute("PRAGMA foreign_keys = ON")
-            connection.execute("PRAGMA synchronous = NORMAL")
-
-        return connection  # type: ignore[no-any-return]
-
-    def _get_thread_connection(self) -> SqliteConnection:
-        """Get or create a connection for the current thread."""
-        try:
-            return cast("SqliteConnection", self._thread_local.connection)
-        except AttributeError:
-            # Connection doesn't exist for this thread yet
-            connection = self._create_connection()
-            self._thread_local.connection = connection
-            return connection
-
-    def _close_thread_connection(self) -> None:
-        """Close the connection for the current thread."""
-        try:
-            connection = self._thread_local.connection
-            connection.close()
-            del self._thread_local.connection
-        except AttributeError:
-            # No connection for this thread
-            pass
-
-    @contextmanager
-    def get_connection(self) -> "Generator[SqliteConnection, None, None]":
-        """Get a thread-local connection.
-
-        Yields:
-            SqliteConnection: A thread-local connection.
-        """
-        yield self._get_thread_connection()
-
-    def close(self) -> None:
-        """Close the thread-local connection if it exists."""
-        self._close_thread_connection()
-
-    def acquire(self) -> SqliteConnection:
-        """Acquire a thread-local connection.
-
-        Returns:
-            SqliteConnection: A thread-local connection
-        """
-        return self._get_thread_connection()
-
-    def release(self, connection: SqliteConnection) -> None:
-        """Release a connection (no-op for thread-local connections).
-
-        Args:
-            connection: The connection to release (ignored)
-        """
-        # No-op: thread-local connections are managed per-thread
-
-    # Compatibility methods that return dummy values
-    def size(self) -> int:
-        """Get pool size (always 1 for thread-local)."""
-        try:
-            _ = self._thread_local.connection
-        except AttributeError:
-            return 0
-        return 1
-
-    def checked_out(self) -> int:
-        """Get number of checked out connections (always 0)."""
-        return 0
+__all__ = ("SqliteConfig", "SqliteConnectionParams")
 
 
 class SqliteConfig(SyncDatabaseConfig[SqliteConnection, SqliteConnectionPool, SqliteDriver]):
@@ -169,7 +58,7 @@ class SqliteConfig(SyncDatabaseConfig[SqliteConnection, SqliteConnectionPool, Sq
         if pool_config is None:
             pool_config = {}
         if "database" not in pool_config or pool_config["database"] == ":memory:":
-            pool_config["database"] = "file::memory:?cache=shared"
+            pool_config["database"] = f"file:memory_{uuid.uuid4().hex}?mode=memory&cache=private"
             pool_config["uri"] = True
 
         super().__init__(