sqlspec 0.11.1__py3-none-any.whl → 0.12.1__py3-none-any.whl

sqlspec/config.py

@@ -0,0 +1,354 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, Optional, TypeVar, Union
+
+from sqlspec.typing import ConnectionT, PoolT  # pyright: ignore
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable
+    from contextlib import AbstractAsyncContextManager, AbstractContextManager
+
+    from sqlglot.dialects.dialect import DialectType
+
+    from sqlspec.driver import AsyncDriverAdapterProtocol, SyncDriverAdapterProtocol
+    from sqlspec.statement.result import StatementResult
+
+
+StatementResultType = Union["StatementResult[dict[str, Any]]", "StatementResult[Any]"]
+
+
+__all__ = (
+    "AsyncConfigT",
+    "AsyncDatabaseConfig",
+    "ConfigT",
+    "DatabaseConfigProtocol",
+    "DriverT",
+    "GenericPoolConfig",
+    "NoPoolAsyncConfig",
+    "NoPoolSyncConfig",
+    "StatementResultType",
+    "SyncConfigT",
+    "SyncDatabaseConfig",
+)
+
+AsyncConfigT = TypeVar("AsyncConfigT", bound="Union[AsyncDatabaseConfig[Any, Any, Any], NoPoolAsyncConfig[Any, Any]]")
+SyncConfigT = TypeVar("SyncConfigT", bound="Union[SyncDatabaseConfig[Any, Any, Any], NoPoolSyncConfig[Any, Any]]")
+ConfigT = TypeVar(
+    "ConfigT",
+    bound="Union[Union[AsyncDatabaseConfig[Any, Any, Any], NoPoolAsyncConfig[Any, Any]], SyncDatabaseConfig[Any, Any, Any], NoPoolSyncConfig[Any, Any]]",
+)
+DriverT = TypeVar("DriverT", bound="Union[SyncDriverAdapterProtocol[Any], AsyncDriverAdapterProtocol[Any]]")
+
+logger = get_logger("config")
+
+
+@dataclass
+class DatabaseConfigProtocol(ABC, Generic[ConnectionT, PoolT, DriverT]):
+    """Protocol defining the interface for database configurations."""
+
+    # Note: __slots__ cannot be used with dataclass fields in Python < 3.10
+    # Concrete subclasses can still use __slots__ for any additional attributes
+    __slots__ = ()
+
+    is_async: "ClassVar[bool]" = field(init=False, default=False)
+    supports_connection_pooling: "ClassVar[bool]" = field(init=False, default=False)
+    supports_native_arrow_import: "ClassVar[bool]" = field(init=False, default=False)
+    supports_native_arrow_export: "ClassVar[bool]" = field(init=False, default=False)
+    supports_native_parquet_import: "ClassVar[bool]" = field(init=False, default=False)
+    supports_native_parquet_export: "ClassVar[bool]" = field(init=False, default=False)
+    connection_type: "type[ConnectionT]" = field(init=False, repr=False, hash=False, compare=False)
+    driver_type: "type[DriverT]" = field(init=False, repr=False, hash=False, compare=False)
+    pool_instance: "Optional[PoolT]" = field(default=None)
+    default_row_type: "type[Any]" = field(init=False)
+    _dialect: "DialectType" = field(default=None, init=False, repr=False, hash=False, compare=False)
+
+    supported_parameter_styles: "ClassVar[tuple[str, ...]]" = ()
+    """Parameter styles supported by this database adapter (e.g., ('qmark', 'named_colon'))."""
+
+    preferred_parameter_style: "ClassVar[str]" = "none"
+    """The preferred/native parameter style for this database."""
+
+    def __hash__(self) -> int:
+        return id(self)
+
+    @property
+    def dialect(self) -> "DialectType":
+        """Get the SQL dialect type lazily.
+
+        This property allows dialect to be set either statically as a class attribute
+        or dynamically via the _get_dialect() method. If a specific adapter needs
+        dynamic dialect detection (e.g., ADBC which supports multiple databases),
+        it can override _get_dialect() to provide custom logic.
+
+        Returns:
+            The SQL dialect type for this database.
+        """
+        if self._dialect is None:
+            self._dialect = self._get_dialect()  # type: ignore[misc]
+        return self._dialect
+
+    def _get_dialect(self) -> "DialectType":
+        """Get the dialect for this database configuration.
+
+        This method should be overridden by configs that need dynamic dialect detection.
+        By default, it looks for the dialect on the driver class.
+
+        Returns:
+            The SQL dialect type.
+        """
+        # Get dialect from driver_class (all drivers must have a dialect attribute)
+        return self.driver_type.dialect
+
+    @abstractmethod
+    def create_connection(self) -> "Union[ConnectionT, Awaitable[ConnectionT]]":
+        """Create and return a new database connection."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def provide_connection(
+        self, *args: Any, **kwargs: Any
+    ) -> "Union[AbstractContextManager[ConnectionT], AbstractAsyncContextManager[ConnectionT]]":
+        """Provide a database connection context manager."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def provide_session(
+        self, *args: Any, **kwargs: Any
+    ) -> "Union[AbstractContextManager[DriverT], AbstractAsyncContextManager[DriverT]]":
+        """Provide a database session context manager."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def connection_config_dict(self) -> "dict[str, Any]":
+        """Return the connection configuration as a dict."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def create_pool(self) -> "Union[PoolT, Awaitable[PoolT]]":
+        """Create and return connection pool."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def close_pool(self) -> "Optional[Awaitable[None]]":
+        """Terminate the connection pool."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def provide_pool(
+        self, *args: Any, **kwargs: Any
+    ) -> "Union[PoolT, Awaitable[PoolT], AbstractContextManager[PoolT], AbstractAsyncContextManager[PoolT]]":
+        """Provide pool instance."""
+        raise NotImplementedError
+
+    def get_signature_namespace(self) -> "dict[str, type[Any]]":
+        """Get the signature namespace for this database configuration.
+
+        This method returns a dictionary of type names to types that should be
+        registered with Litestar's signature namespace to prevent serialization
+        attempts on database-specific types.
+
+        Returns:
+            Dictionary mapping type names to types.
+        """
+        namespace: dict[str, type[Any]] = {}
+
+        # Add the driver and config types
+        if hasattr(self, "driver_type") and self.driver_type:
+            namespace[self.driver_type.__name__] = self.driver_type
+
+        namespace[self.__class__.__name__] = self.__class__
+
+        # Add connection type(s)
+        if hasattr(self, "connection_type") and self.connection_type:
+            connection_type = self.connection_type
+
+            # Handle Union types (like AsyncPG's Union[Connection, PoolConnectionProxy])
+            if hasattr(connection_type, "__args__"):
+                # It's a generic type, extract the actual types
+                for arg_type in connection_type.__args__:  # type: ignore[attr-defined]
+                    if arg_type and hasattr(arg_type, "__name__"):
+                        namespace[arg_type.__name__] = arg_type
+            elif hasattr(connection_type, "__name__"):
+                # Regular type
+                namespace[connection_type.__name__] = connection_type
+
+        return namespace
+
+
+@dataclass
+class NoPoolSyncConfig(DatabaseConfigProtocol[ConnectionT, None, DriverT]):
+    """Base class for a sync database configurations that do not implement a pool."""
+
+    __slots__ = ()
+
+    is_async: "ClassVar[bool]" = field(init=False, default=False)
+    supports_connection_pooling: "ClassVar[bool]" = field(init=False, default=False)
+    pool_instance: None = None
+
+    def create_connection(self) -> ConnectionT:
+        """Create connection with instrumentation."""
+        raise NotImplementedError
+
+    def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractContextManager[ConnectionT]":
+        """Provide connection with instrumentation."""
+        raise NotImplementedError
+
+    def provide_session(self, *args: Any, **kwargs: Any) -> "AbstractContextManager[DriverT]":
+        """Provide session with instrumentation."""
+        raise NotImplementedError
+
+    def create_pool(self) -> None:
+        return None
+
+    def close_pool(self) -> None:
+        return None
+
+    def provide_pool(self, *args: Any, **kwargs: Any) -> None:
+        return None
+
+
+@dataclass
+class NoPoolAsyncConfig(DatabaseConfigProtocol[ConnectionT, None, DriverT]):
+    """Base class for an async database configurations that do not implement a pool."""
+
+    __slots__ = ()
+
+    is_async: "ClassVar[bool]" = field(init=False, default=True)
+    supports_connection_pooling: "ClassVar[bool]" = field(init=False, default=False)
+    pool_instance: None = None
+
+    async def create_connection(self) -> ConnectionT:
+        """Create connection with instrumentation."""
+        raise NotImplementedError
+
+    def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractAsyncContextManager[ConnectionT]":
+        """Provide connection with instrumentation."""
+        raise NotImplementedError
+
+    def provide_session(self, *args: Any, **kwargs: Any) -> "AbstractAsyncContextManager[DriverT]":
+        """Provide session with instrumentation."""
+        raise NotImplementedError
+
+    async def create_pool(self) -> None:
+        return None
+
+    async def close_pool(self) -> None:
+        return None
+
+    def provide_pool(self, *args: Any, **kwargs: Any) -> None:
+        return None
+
+
+@dataclass
+class GenericPoolConfig:
+    """Generic Database Pool Configuration."""
+
+    __slots__ = ()
+
+
+@dataclass
+class SyncDatabaseConfig(DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]):
+    """Generic Sync Database Configuration."""
+
+    __slots__ = ()
+
+    is_async: "ClassVar[bool]" = field(init=False, default=False)
+    supports_connection_pooling: "ClassVar[bool]" = field(init=False, default=True)
+
+    def create_pool(self) -> PoolT:
+        """Create pool with instrumentation.
+
+        Returns:
+            The created pool.
+        """
+        if self.pool_instance is not None:
+            return self.pool_instance
+        self.pool_instance = self._create_pool()  # type: ignore[misc]
+        return self.pool_instance
+
+    def close_pool(self) -> None:
+        """Close pool with instrumentation."""
+        self._close_pool()
+
+    def provide_pool(self, *args: Any, **kwargs: Any) -> PoolT:
+        """Provide pool instance."""
+        if self.pool_instance is None:
+            self.pool_instance = self.create_pool()  # type: ignore[misc]
+        return self.pool_instance
+
+    def create_connection(self) -> ConnectionT:
+        """Create connection with instrumentation."""
+        raise NotImplementedError
+
+    def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractContextManager[ConnectionT]":
+        """Provide connection with instrumentation."""
+        raise NotImplementedError
+
+    def provide_session(self, *args: Any, **kwargs: Any) -> "AbstractContextManager[DriverT]":
+        """Provide session with instrumentation."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _create_pool(self) -> PoolT:
+        """Actual pool creation implementation."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _close_pool(self) -> None:
+        """Actual pool destruction implementation."""
+        raise NotImplementedError
+
+
+@dataclass
+class AsyncDatabaseConfig(DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]):
+    """Generic Async Database Configuration."""
+
+    __slots__ = ()
+
+    is_async: "ClassVar[bool]" = field(init=False, default=True)
+    supports_connection_pooling: "ClassVar[bool]" = field(init=False, default=True)
+
+    async def create_pool(self) -> PoolT:
+        """Create pool with instrumentation.
+
+        Returns:
+            The created pool.
+        """
+        if self.pool_instance is not None:
+            return self.pool_instance
+        self.pool_instance = await self._create_pool()  # type: ignore[misc]
+        return self.pool_instance
+
+    async def close_pool(self) -> None:
+        """Close pool with instrumentation."""
+        await self._close_pool()
+
+    async def provide_pool(self, *args: Any, **kwargs: Any) -> PoolT:
+        """Provide pool instance."""
+        if self.pool_instance is None:
+            self.pool_instance = await self.create_pool()  # type: ignore[misc]
+        return self.pool_instance
+
+    async def create_connection(self) -> ConnectionT:
+        """Create connection with instrumentation."""
+        raise NotImplementedError
+
+    def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractAsyncContextManager[ConnectionT]":
+        """Provide connection with instrumentation."""
+        raise NotImplementedError
+
+    def provide_session(self, *args: Any, **kwargs: Any) -> "AbstractAsyncContextManager[DriverT]":
+        """Provide session with instrumentation."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def _create_pool(self) -> PoolT:
+        """Actual async pool creation implementation."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def _close_pool(self) -> None:
+        """Actual async pool destruction implementation."""
+        raise NotImplementedError

sqlspec/driver/__init__.py

@@ -0,0 +1,22 @@
+"""Driver protocols and base classes for database adapters."""
+
+from typing import Union
+
+from sqlspec.driver import mixins
+from sqlspec.driver._async import AsyncDriverAdapterProtocol
+from sqlspec.driver._common import CommonDriverAttributesMixin
+from sqlspec.driver._sync import SyncDriverAdapterProtocol
+from sqlspec.typing import ConnectionT, RowT
+
+__all__ = (
+    "AsyncDriverAdapterProtocol",
+    "CommonDriverAttributesMixin",
+    "DriverAdapterProtocol",
+    "SyncDriverAdapterProtocol",
+    "mixins",
+)
+
+# Type alias for convenience
+DriverAdapterProtocol = Union[
+    SyncDriverAdapterProtocol[ConnectionT, RowT], AsyncDriverAdapterProtocol[ConnectionT, RowT]
+]

sqlspec/driver/_async.py

@@ -0,0 +1,252 @@
+"""Asynchronous driver protocol implementation."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
+
+from sqlspec.driver._common import CommonDriverAttributesMixin
+from sqlspec.statement.builder import DeleteBuilder, InsertBuilder, QueryBuilder, SelectBuilder, UpdateBuilder
+from sqlspec.statement.filters import StatementFilter
+from sqlspec.statement.result import SQLResult
+from sqlspec.statement.sql import SQL, SQLConfig, Statement
+from sqlspec.typing import ConnectionT, DictRow, ModelDTOT, RowT, StatementParameters
+
+if TYPE_CHECKING:
+    from sqlspec.statement.result import DMLResultDict, ScriptResultDict, SelectResultDict
+
+__all__ = ("AsyncDriverAdapterProtocol",)
+
+
+EMPTY_FILTERS: "list[StatementFilter]" = []
+
+
+class AsyncDriverAdapterProtocol(CommonDriverAttributesMixin[ConnectionT, RowT], ABC):
+    __slots__ = ()
+
+    def __init__(
+        self,
+        connection: "ConnectionT",
+        config: "Optional[SQLConfig]" = None,
+        default_row_type: "type[DictRow]" = DictRow,
+    ) -> None:
+        """Initialize async driver adapter.
+
+        Args:
+            connection: The database connection
+            config: SQL statement configuration
+            default_row_type: Default row type for results (DictRow, TupleRow, etc.)
+        """
+        super().__init__(connection=connection, config=config, default_row_type=default_row_type)
+
+    def _build_statement(
+        self,
+        statement: "Union[Statement, QueryBuilder[Any]]",
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        _config: "Optional[SQLConfig]" = None,
+        **kwargs: Any,
+    ) -> "SQL":
+        # Use driver's config if none provided
+        _config = _config or self.config
+
+        if isinstance(statement, QueryBuilder):
+            return statement.to_statement(config=_config)
+        # If statement is already a SQL object, return it as-is
+        if isinstance(statement, SQL):
+            return statement
+        return SQL(statement, *parameters, _dialect=self.dialect, _config=_config, **kwargs)
+
+    @abstractmethod
+    async def _execute_statement(
+        self, statement: "SQL", connection: "Optional[ConnectionT]" = None, **kwargs: Any
+    ) -> "Union[SelectResultDict, DMLResultDict, ScriptResultDict]":
+        """Actual execution implementation by concrete drivers, using the raw connection.
+
+        Returns one of the standardized result dictionaries based on the statement type.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def _wrap_select_result(
+        self,
+        statement: "SQL",
+        result: "SelectResultDict",
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "Union[SQLResult[ModelDTOT], SQLResult[RowT]]":
+        raise NotImplementedError
+
+    @abstractmethod
+    async def _wrap_execute_result(
+        self, statement: "SQL", result: "Union[DMLResultDict, ScriptResultDict]", **kwargs: Any
+    ) -> "SQLResult[RowT]":
+        raise NotImplementedError
+
+    # Type-safe overloads based on the refactor plan pattern
+    @overload
+    async def execute(
+        self,
+        statement: "SelectBuilder",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: "type[ModelDTOT]",
+        _connection: "Optional[ConnectionT]" = None,
+        _config: "Optional[SQLConfig]" = None,
+        **kwargs: Any,
+    ) -> "SQLResult[ModelDTOT]": ...
+
+    @overload
+    async def execute(
+        self,
+        statement: "SelectBuilder",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: None = None,
+        _connection: "Optional[ConnectionT]" = None,
+        _config: "Optional[SQLConfig]" = None,
+        **kwargs: Any,
+    ) -> "SQLResult[RowT]": ...
+
+    @overload
+    async def execute(
+        self,
+        statement: "Union[InsertBuilder, UpdateBuilder, DeleteBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        _connection: "Optional[ConnectionT]" = None,
+        _config: "Optional[SQLConfig]" = None,
+        **kwargs: Any,
+    ) -> "SQLResult[RowT]": ...
+
+    @overload
+    async def execute(
+        self,
+        statement: "Union[str, SQL]",  # exp.Expression
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: "type[ModelDTOT]",
+        _connection: "Optional[ConnectionT]" = None,
+        _config: "Optional[SQLConfig]" = None,
+        **kwargs: Any,
+    ) -> "SQLResult[ModelDTOT]": ...
+
+    @overload
+    async def execute(
+        self,
+        statement: "Union[str, SQL]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: None = None,
+        _connection: "Optional[ConnectionT]" = None,
+        _config: "Optional[SQLConfig]" = None,
+        **kwargs: Any,
+    ) -> "SQLResult[RowT]": ...
+
+    async def execute(
+        self,
+        statement: "Union[SQL, Statement, QueryBuilder[Any]]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        _connection: "Optional[ConnectionT]" = None,
+        _config: "Optional[SQLConfig]" = None,
+        **kwargs: Any,
+    ) -> "Union[SQLResult[ModelDTOT], SQLResult[RowT]]":
+        sql_statement = self._build_statement(statement, *parameters, _config=_config or self.config, **kwargs)
+        result = await self._execute_statement(
+            statement=sql_statement, connection=self._connection(_connection), **kwargs
+        )
+
+        if self.returns_rows(sql_statement.expression):
+            return await self._wrap_select_result(
+                sql_statement, cast("SelectResultDict", result), schema_type=schema_type, **kwargs
+            )
+        return await self._wrap_execute_result(
+            sql_statement, cast("Union[DMLResultDict, ScriptResultDict]", result), **kwargs
+        )
+
+    async def execute_many(
+        self,
+        statement: "Union[SQL, Statement, QueryBuilder[Any]]",  # QueryBuilder for DMLs will likely not return rows.
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        _connection: "Optional[ConnectionT]" = None,
+        _config: "Optional[SQLConfig]" = None,
+        **kwargs: Any,
+    ) -> "SQLResult[RowT]":
+        # Separate parameters from filters
+        param_sequences = []
+        filters = []
+        for param in parameters:
+            if isinstance(param, StatementFilter):
+                filters.append(param)
+            else:
+                param_sequences.append(param)
+
+        # Use first parameter as the sequence for execute_many
+        param_sequence = param_sequences[0] if param_sequences else None
+        # Convert tuple to list if needed
+        if isinstance(param_sequence, tuple):
+            param_sequence = list(param_sequence)
+        # Ensure param_sequence is a list or None
+        if param_sequence is not None and not isinstance(param_sequence, list):
+            param_sequence = list(param_sequence) if hasattr(param_sequence, "__iter__") else None
+        sql_statement = self._build_statement(statement, _config=_config or self.config, **kwargs)
+        sql_statement = sql_statement.as_many(param_sequence)
+        result = await self._execute_statement(
+            statement=sql_statement,
+            connection=self._connection(_connection),
+            parameters=param_sequence,
+            is_many=True,
+            **kwargs,
+        )
+        return await self._wrap_execute_result(
+            sql_statement, cast("Union[DMLResultDict, ScriptResultDict]", result), **kwargs
+        )
+
+    async def execute_script(
+        self,
+        statement: "Union[str, SQL]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        _connection: "Optional[ConnectionT]" = None,
+        _config: "Optional[SQLConfig]" = None,
+        **kwargs: Any,
+    ) -> "SQLResult[RowT]":
+        param_values = []
+        filters = []
+        for param in parameters:
+            if isinstance(param, StatementFilter):
+                filters.append(param)
+            else:
+                param_values.append(param)
+
+        # Use first parameter as the primary parameter value, or None if no parameters
+        primary_params = param_values[0] if param_values else None
+
+        script_config = _config or self.config
+        if script_config.enable_validation:
+            script_config = SQLConfig(
+                enable_parsing=script_config.enable_parsing,
+                enable_validation=False,
+                enable_transformations=script_config.enable_transformations,
+                enable_analysis=script_config.enable_analysis,
+                strict_mode=False,
+                cache_parsed_expression=script_config.cache_parsed_expression,
+                parameter_converter=script_config.parameter_converter,
+                parameter_validator=script_config.parameter_validator,
+                analysis_cache_size=script_config.analysis_cache_size,
+                allowed_parameter_styles=script_config.allowed_parameter_styles,
+                target_parameter_style=script_config.target_parameter_style,
+                allow_mixed_parameter_styles=script_config.allow_mixed_parameter_styles,
+            )
+        sql_statement = SQL(statement, primary_params, *filters, _dialect=self.dialect, _config=script_config, **kwargs)
+        sql_statement = sql_statement.as_script()
+        script_output = await self._execute_statement(
+            statement=sql_statement, connection=self._connection(_connection), is_script=True, **kwargs
+        )
+        if isinstance(script_output, str):
+            result = SQLResult[RowT](statement=sql_statement, data=[], operation_type="SCRIPT")
+            result.total_statements = 1
+            result.successful_statements = 1
+            return result
+        # Wrap the ScriptResultDict using the driver's wrapper
+        return await self._wrap_execute_result(sql_statement, cast("ScriptResultDict", script_output), **kwargs)