sqlspec 0.7.1__py3-none-any.whl → 0.9.0__py3-none-any.whl

sqlspec/__init__.py

@@ -0,0 +1,15 @@
+from sqlspec import adapters, base, exceptions, extensions, filters, typing, utils
+from sqlspec.__metadata__ import __version__
+from sqlspec.base import SQLSpec
+
+__all__ = (
+    "SQLSpec",
+    "__version__",
+    "adapters",
+    "base",
+    "exceptions",
+    "extensions",
+    "filters",
+    "typing",
+    "utils",
+)

sqlspec/_serialization.py

@@ -60,12 +60,26 @@ __all__ = (
 
 
 def convert_datetime_to_gmt_iso(dt: datetime.datetime) -> str:  # pragma: no cover
-    """Handle datetime serialization for nested timestamps."""
+    """Handle datetime serialization for nested timestamps.
+
+    Args:
+        dt: The datetime to convert.
+
+    Returns:
+        The ISO formatted datetime string.
+    """
     if not dt.tzinfo:
         dt = dt.replace(tzinfo=datetime.timezone.utc)
     return dt.isoformat().replace("+00:00", "Z")
 
 
 def convert_date_to_iso(dt: datetime.date) -> str:  # pragma: no cover
-    """Handle datetime serialization for nested timestamps."""
+    """Handle datetime serialization for nested timestamps.
+
+    Args:
+        dt: The date to convert.
+
+    Returns:
+        The ISO formatted date string.
+    """
     return dt.isoformat()

sqlspec/_typing.py

@@ -1,8 +1,10 @@
+# ruff: noqa: RUF100, PLR0913, A002, DOC201, PLR6301
 """This is a simple wrapper around a few important classes in each library.
 
 This is used to ensure compatibility when one or more of the libraries are installed.
 """
 
+from collections.abc import Iterable, Mapping
 from enum import Enum
 from typing import (
     Any,
@@ -29,7 +31,7 @@ T_co = TypeVar("T_co", covariant=True)
 
 try:
     from pydantic import (
-        BaseModel,
+        BaseModel,  # pyright: ignore[reportAssignmentType]
         FailFast,  # pyright: ignore[reportGeneralTypeIssues,reportAssignmentType]
         TypeAdapter,
     )
@@ -96,7 +98,7 @@ except ImportError:
 
         def validate_python(
             self,
-            object: Any,  # noqa: A002
+            object: Any,
             /,
             *,
             strict: "Optional[bool]" = None,
@@ -127,10 +129,7 @@ try:
 except ImportError:
     import enum
     from collections.abc import Iterable
-    from typing import TYPE_CHECKING, Callable, Optional, Union
-
-    if TYPE_CHECKING:
-        from collections.abc import Iterable
+    from typing import Callable, Optional, Union
 
     @dataclass_transform()
     @runtime_checkable
@@ -174,7 +173,6 @@ except ImportError:
             """Placeholder init"""
 
         def create_instance(self, **kwargs: Any) -> "T":
-            """Placeholder implementation"""
             return cast("T", kwargs)
 
         def update_instance(self, instance: "T", **kwargs: Any) -> "T":
@@ -198,11 +196,46 @@ EmptyType = Union[Literal[EmptyEnum.EMPTY], UnsetType]
 Empty: Final = EmptyEnum.EMPTY
 
 
+try:
+    from pyarrow import Table as ArrowTable
+
+    PYARROW_INSTALLED = True
+except ImportError:
+
+    @runtime_checkable
+    class ArrowTable(Protocol):  # type: ignore[no-redef]
+        """Placeholder Implementation"""
+
+        def to_batches(self, batch_size: int) -> Any: ...
+        def num_rows(self) -> int: ...
+        def num_columns(self) -> int: ...
+        def to_pydict(self) -> dict[str, Any]: ...
+        def to_string(self) -> str: ...
+        def from_arrays(
+            self,
+            arrays: list[Any],
+            names: "Optional[list[str]]" = None,
+            schema: "Optional[Any]" = None,
+            metadata: "Optional[Mapping[str, Any]]" = None,
+        ) -> Any: ...
+        def from_pydict(
+            self,
+            mapping: dict[str, Any],
+            schema: "Optional[Any]" = None,
+            metadata: "Optional[Mapping[str, Any]]" = None,
+        ) -> Any: ...
+        def from_batches(self, batches: Iterable[Any], schema: Optional[Any] = None) -> Any: ...
+
+    PYARROW_INSTALLED = False  # pyright: ignore[reportConstantRedefinition]
+
+
 __all__ = (
     "LITESTAR_INSTALLED",
     "MSGSPEC_INSTALLED",
+    "PYARROW_INSTALLED",
     "PYDANTIC_INSTALLED",
     "UNSET",
+    "ArrowTable",
     "BaseModel",
     "DTOData",
     "DataclassProtocol",

sqlspec/adapters/adbc/__init__.py

@@ -0,0 +1,7 @@
+from sqlspec.adapters.adbc.config import AdbcConfig
+from sqlspec.adapters.adbc.driver import AdbcDriver
+
+__all__ = (
+    "AdbcConfig",
+    "AdbcDriver",
+)

sqlspec/adapters/adbc/config.py

@@ -1,21 +1,24 @@
 from contextlib import contextmanager
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Optional, Union
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast
 
+from adbc_driver_manager.dbapi import Connection
+
+from sqlspec.adapters.adbc.driver import AdbcDriver
 from sqlspec.base import NoPoolSyncConfig
+from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import Empty, EmptyType
+from sqlspec.utils.module_loader import import_string
 
 if TYPE_CHECKING:
     from collections.abc import Generator
-    from typing import Any
 
-    from adbc_driver_manager.dbapi import Connection
 
-__all__ = ("AdbcDatabaseConfig",)
+__all__ = ("AdbcConfig",)
 
 
 @dataclass
-class AdbcDatabaseConfig(NoPoolSyncConfig["Connection"]):
+class AdbcConfig(NoPoolSyncConfig["Connection", "AdbcDriver"]):
     """Configuration for ADBC connections.
 
     This class provides configuration options for ADBC database connections using the
@@ -25,23 +28,186 @@ class AdbcDatabaseConfig(NoPoolSyncConfig["Connection"]):
     uri: "Union[str, EmptyType]" = Empty
     """Database URI"""
     driver_name: "Union[str, EmptyType]" = Empty
-    """Name of the ADBC driver to use"""
+    """Full dotted path to the ADBC driver's connect function (e.g., 'adbc_driver_sqlite.dbapi.connect')"""
     db_kwargs: "Optional[dict[str, Any]]" = None
     """Additional database-specific connection parameters"""
+    conn_kwargs: "Optional[dict[str, Any]]" = None
+    """Additional database-specific connection parameters"""
+    connection_type: "type[Connection]" = field(init=False, default_factory=lambda: Connection)
+    """Type of the connection object"""
+    driver_type: "type[AdbcDriver]" = field(init=False, default_factory=lambda: AdbcDriver)  # type: ignore[type-abstract,unused-ignore]
+    """Type of the driver object"""
+    pool_instance: None = field(init=False, default=None)
+    """No connection pool is used for ADBC connections"""
+    _is_in_memory: bool = field(init=False, default=False)
+    """Flag indicating if the connection is for an in-memory database"""
+
+    def _set_adbc(self) -> str:  # noqa: PLR0912
+        """Identify the driver type based on the URI (if provided) or preset driver name.
+
+        Also sets the `_is_in_memory` flag for specific in-memory URIs.
+
+        Raises:
+            ImproperConfigurationError: If the driver name is not recognized or supported.
+
+        Returns:
+            str: The driver name to be used for the connection.
+        """
+
+        if isinstance(self.driver_name, str):
+            if self.driver_name != "adbc_driver_sqlite.dbapi.connect" and self.driver_name in {
+                "sqlite",
+                "sqlite3",
+                "adbc_driver_sqlite",
+            }:
+                self.driver_name = "adbc_driver_sqlite.dbapi.connect"
+            elif self.driver_name != "adbc_driver_duckdb.dbapi.connect" and self.driver_name in {
+                "duckdb",
+                "adbc_driver_duckdb",
+            }:
+                self.driver_name = "adbc_driver_duckdb.dbapi.connect"
+            elif self.driver_name != "adbc_driver_postgresql.dbapi.connect" and self.driver_name in {
+                "postgres",
+                "adbc_driver_postgresql",
+                "postgresql",
+                "pg",
+            }:
+                self.driver_name = "adbc_driver_postgresql.dbapi.connect"
+            elif self.driver_name != "adbc_driver_snowflake.dbapi.connect" and self.driver_name in {
+                "snowflake",
+                "adbc_driver_snowflake",
+                "sf",
+            }:
+                self.driver_name = "adbc_driver_snowflake.dbapi.connect"
+            elif self.driver_name != "adbc_driver_bigquery.dbapi.connect" and self.driver_name in {
+                "bigquery",
+                "adbc_driver_bigquery",
+                "bq",
+            }:
+                self.driver_name = "adbc_driver_bigquery.dbapi.connect"
+            elif self.driver_name != "adbc_driver_flightsql.dbapi.connect" and self.driver_name in {
+                "flightsql",
+                "adbc_driver_flightsql",
+                "grpc",
+            }:
+                self.driver_name = "adbc_driver_flightsql.dbapi.connect"
+            return self.driver_name
+
+        # If driver_name wasn't explicit, try to determine from URI
+        if isinstance(self.uri, str) and self.uri.startswith("postgresql://"):
+            self.driver_name = "adbc_driver_postgresql.dbapi.connect"
+        elif isinstance(self.uri, str) and self.uri.startswith("sqlite://"):
+            self.driver_name = "adbc_driver_sqlite.dbapi.connect"
+        elif isinstance(self.uri, str) and self.uri.startswith("grpc://"):
+            self.driver_name = "adbc_driver_flightsql.dbapi.connect"
+        elif isinstance(self.uri, str) and self.uri.startswith("snowflake://"):
+            self.driver_name = "adbc_driver_snowflake.dbapi.connect"
+        elif isinstance(self.uri, str) and self.uri.startswith("bigquery://"):
+            self.driver_name = "adbc_driver_bigquery.dbapi.connect"
+        elif isinstance(self.uri, str) and self.uri.startswith("duckdb://"):
+            self.driver_name = "adbc_driver_duckdb.dbapi.connect"
+
+        # Check if we successfully determined a driver name
+        if self.driver_name is Empty or not isinstance(self.driver_name, str):
+            msg = (
+                "Could not determine ADBC driver connect path. Please specify 'driver_name' "
+                "(e.g., 'adbc_driver_sqlite.dbapi.connect') or provide a supported 'uri'. "
+                f"URI: {self.uri}, Driver Name: {self.driver_name}"
+            )
+            raise ImproperConfigurationError(msg)
+        return self.driver_name
 
     @property
-    def connection_params(self) -> "dict[str, Any]":
-        """Return the connection parameters as a dict."""
-        return {
-            k: v
-            for k, v in {"uri": self.uri, "driver": self.driver_name, **(self.db_kwargs or {})}.items()
-            if v is not Empty
-        }
+    def connection_config_dict(self) -> "dict[str, Any]":
+        """Return the connection configuration as a dict.
+
+        Omits the 'uri' key for known in-memory database types.
+
+        Returns:
+            A string keyed dict of config kwargs for the adbc_driver_manager.dbapi.connect function.
+        """
+        config = {}
+        db_kwargs = self.db_kwargs or {}
+        conn_kwargs = self.conn_kwargs or {}
+        if isinstance(self.uri, str) and self.uri.startswith("sqlite://"):
+            db_kwargs["uri"] = self.uri.replace("sqlite://", "")
+        elif isinstance(self.uri, str) and self.uri.startswith("duckdb://"):
+            db_kwargs["path"] = self.uri.replace("duckdb://", "")
+        elif isinstance(self.uri, str):
+            db_kwargs["uri"] = self.uri
+        if isinstance(self.driver_name, str) and self.driver_name.startswith("adbc_driver_bigquery"):
+            config["db_kwargs"] = db_kwargs
+        else:
+            config = db_kwargs
+        if conn_kwargs:
+            config["conn_kwargs"] = conn_kwargs
+        return config
+
+    def _get_connect_func(self) -> "Callable[..., Connection]":
+        self._set_adbc()
+        driver_path = cast("str", self.driver_name)
+        try:
+            connect_func = import_string(driver_path)
+        except ImportError as e:
+            # Check if the error is likely due to missing suffix and try again
+            if ".dbapi.connect" not in driver_path:
+                try:
+                    driver_path += ".dbapi.connect"
+                    connect_func = import_string(driver_path)
+                except ImportError as e2:
+                    msg = f"Failed to import ADBC connect function from '{self.driver_name}' or '{driver_path}'. Is the driver installed and the path correct? Original error: {e} / {e2}"
+                    raise ImproperConfigurationError(msg) from e2
+            else:
+                # Original import failed, and suffix was already present or added
+                msg = f"Failed to import ADBC connect function from '{driver_path}'. Is the driver installed and the path correct? Original error: {e}"
+                raise ImproperConfigurationError(msg) from e
+        if not callable(connect_func):
+            msg = f"The path '{driver_path}' did not resolve to a callable function."
+            raise ImproperConfigurationError(msg)
+        return connect_func  # type: ignore[no-any-return]
+
+    def create_connection(self) -> "Connection":
+        """Create and return a new database connection using the specific driver.
+
+        Returns:
+            A new ADBC connection instance.
+
+        Raises:
+            ImproperConfigurationError: If the connection could not be established.
+        """
+        try:
+            connect_func = self._get_connect_func()
+            return connect_func(**self.connection_config_dict)
+        except Exception as e:
+            # Include driver name in error message for better context
+            driver_name = self.driver_name if isinstance(self.driver_name, str) else "Unknown/Missing"
+            # Use the potentially modified driver_path from _get_connect_func if available,
+            # otherwise fallback to self.driver_name for the error message.
+            # This requires _get_connect_func to potentially return the used path or store it.
+            # For simplicity now, we stick to self.driver_name in the message.
+            msg = f"Could not configure the ADBC connection using driver path '{driver_name}'. Error: {e!s}"
+            raise ImproperConfigurationError(msg) from e
 
     @contextmanager
     def provide_connection(self, *args: "Any", **kwargs: "Any") -> "Generator[Connection, None, None]":
-        """Create and provide a database connection."""
-        from adbc_driver_manager.dbapi import connect
+        """Create and provide a database connection using the specific driver.
 
-        with connect(**self.connection_params) as connection:
+        Yields:
+            Connection: A database connection instance.
+        """
+
+        connection = self.create_connection()
+        try:
             yield connection
+        finally:
+            connection.close()
+
+    @contextmanager
+    def provide_session(self, *args: Any, **kwargs: Any) -> "Generator[AdbcDriver, None, None]":
+        """Create and provide a database session.
+
+        Yields:
+            An ADBC driver instance with an active connection.
+        """
+        with self.provide_connection(*args, **kwargs) as connection:
+            yield self.driver_type(connection)