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

sqlspec/_typing.py

@@ -1,4 +1,4 @@
-# ruff: noqa: RUF100, PLR0913, A002, DOC201, PLR6301
+# ruff: noqa: RUF100, PLR0913, A002, DOC201, PLR6301, PLR0917, ARG004
 """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.
@@ -6,6 +6,7 @@ This is used to ensure compatibility when one or more of the libraries are insta
 
 from collections.abc import Iterable, Mapping
 from enum import Enum
+from importlib.util import find_spec
 from typing import Any, ClassVar, Final, Optional, Protocol, Union, cast, runtime_checkable
 
 from typing_extensions import Literal, TypeVar, dataclass_transform
@@ -13,9 +14,13 @@ from typing_extensions import Literal, TypeVar, dataclass_transform
 
 @runtime_checkable
 class DataclassProtocol(Protocol):
-    """Protocol for instance checking dataclasses."""
+    """Protocol for instance checking dataclasses.
 
-    __dataclass_fields__: ClassVar[dict[str, Any]]
+    This protocol only requires the presence of `__dataclass_fields__`, which is the
+    standard attribute that Python's dataclasses module adds to all dataclass instances.
+    """
+
+    __dataclass_fields__: "ClassVar[dict[str, Any]]"
 
 
 T = TypeVar("T")
@@ -32,19 +37,18 @@ try:
 except ImportError:
     from dataclasses import dataclass
 
-    @runtime_checkable
     class BaseModel(Protocol):  # type: ignore[no-redef]
         """Placeholder Implementation"""
 
-        model_fields: ClassVar[dict[str, Any]]
+        model_fields: "ClassVar[dict[str, Any]]"
 
         def model_dump(
             self,
             /,
             *,
-            include: Optional[Any] = None,
-            exclude: Optional[Any] = None,
-            context: Optional[Any] = None,
+            include: "Optional[Any]" = None,
+            exclude: "Optional[Any]" = None,
+            context: "Optional[Any]" = None,
             by_alias: bool = False,
             exclude_unset: bool = False,
             exclude_defaults: bool = False,
@@ -60,9 +64,9 @@ except ImportError:
             self,
             /,
             *,
-            include: Optional[Any] = None,
-            exclude: Optional[Any] = None,
-            context: Optional[Any] = None,
+            include: "Optional[Any]" = None,
+            exclude: "Optional[Any]" = None,
+            context: "Optional[Any]" = None,
             by_alias: bool = False,
             exclude_unset: bool = False,
             exclude_defaults: bool = False,
@@ -82,9 +86,9 @@ except ImportError:
             self,
             type: Any,  # noqa: A002
             *,
-            config: Optional[Any] = None,
+            config: "Optional[Any]" = None,
             _parent_depth: int = 2,
-            module: Optional[str] = None,
+            module: "Optional[str]" = None,
         ) -> None:
             """Init"""
 
@@ -93,10 +97,10 @@ except ImportError:
             object: Any,
             /,
             *,
-            strict: Optional[bool] = None,
-            from_attributes: Optional[bool] = None,
-            context: Optional[dict[str, Any]] = None,
-            experimental_allow_partial: Union[bool, Literal["off", "on", "trailing-strings"]] = False,
+            strict: "Optional[bool]" = None,
+            from_attributes: "Optional[bool]" = None,
+            context: "Optional[dict[str, Any]]" = None,
+            experimental_allow_partial: "Union[bool, Literal['off', 'on', 'trailing-strings']]" = False,
         ) -> "T_co":
             """Stub"""
             return cast("T_co", object)
@@ -128,16 +132,16 @@ except ImportError:
     class Struct(Protocol):  # type: ignore[no-redef]
         """Placeholder Implementation"""
 
-        __struct_fields__: ClassVar[tuple[str, ...]]
+        __struct_fields__: "ClassVar[tuple[str, ...]]"
 
     def convert(  # type: ignore[no-redef]
         obj: Any,
-        type: Union[Any, type[T]],  # noqa: A002
+        type: "Union[Any, type[T]]",  # noqa: A002
         *,
         strict: bool = True,
         from_attributes: bool = False,
-        dec_hook: Optional[Callable[[type, Any], Any]] = None,
-        builtin_types: Optional[Iterable[type]] = None,
+        dec_hook: "Optional[Callable[[type, Any], Any]]" = None,
+        builtin_types: "Optional[Iterable[type]]" = None,
         str_keys: bool = False,
     ) -> "Union[T, Any]":
         """Placeholder implementation"""
@@ -188,55 +192,353 @@ EmptyType = Union[Literal[EmptyEnum.EMPTY], UnsetType]
 Empty: Final = EmptyEnum.EMPTY
 
 
+@runtime_checkable
+class ArrowTableResult(Protocol):
+    """This is a typed shim for pyarrow.Table."""
+
+    def to_batches(self, batch_size: int) -> Any:
+        return None
+
+    @property
+    def num_rows(self) -> int:
+        return 0
+
+    @property
+    def num_columns(self) -> int:
+        return 0
+
+    def to_pydict(self) -> dict[str, Any]:
+        return {}
+
+    def to_string(self) -> str:
+        return ""
+
+    def from_arrays(
+        self,
+        arrays: list[Any],
+        names: "Optional[list[str]]" = None,
+        schema: "Optional[Any]" = None,
+        metadata: "Optional[Mapping[str, Any]]" = None,
+    ) -> Any:
+        return None
+
+    def from_pydict(
+        self, mapping: dict[str, Any], schema: "Optional[Any]" = None, metadata: "Optional[Mapping[str, Any]]" = None
+    ) -> Any:
+        return None
+
+    def from_batches(self, batches: Iterable[Any], schema: Optional[Any] = None) -> Any:
+        return None
+
+
+@runtime_checkable
+class ArrowRecordBatchResult(Protocol):
+    """This is a typed shim for pyarrow.RecordBatch."""
+
+    def num_rows(self) -> int:
+        return 0
+
+    def num_columns(self) -> int:
+        return 0
+
+    def to_pydict(self) -> dict[str, Any]:
+        return {}
+
+    def to_pandas(self) -> Any:
+        return None
+
+    def schema(self) -> Any:
+        return None
+
+    def column(self, i: int) -> Any:
+        return None
+
+    def slice(self, offset: int = 0, length: "Optional[int]" = None) -> Any:
+        return None
+
+
 try:
+    from pyarrow import RecordBatch as ArrowRecordBatch
     from pyarrow import Table as ArrowTable
 
     PYARROW_INSTALLED = True
 except ImportError:
+    ArrowTable = ArrowTableResult  # type: ignore[assignment,misc]
+    ArrowRecordBatch = ArrowRecordBatchResult  # type: ignore[assignment,misc]
 
-    @runtime_checkable
-    class ArrowTable(Protocol):  # type: ignore[no-redef]
-        """Placeholder Implementation"""
+    PYARROW_INSTALLED = False  # pyright: ignore[reportConstantRedefinition]
+
+
+try:
+    from opentelemetry import trace  # pyright: ignore[reportMissingImports, reportAssignmentType]
+    from opentelemetry.trace import (  # pyright: ignore[reportMissingImports, reportAssignmentType]
+        Span,  # pyright: ignore[reportMissingImports, reportAssignmentType]
+        Status,
+        StatusCode,
+        Tracer,  # pyright: ignore[reportMissingImports, reportAssignmentType]
+    )
+
+    OPENTELEMETRY_INSTALLED = True
+except ImportError:
+    # Define shims for when opentelemetry is not installed
+
+    class Span:  # type: ignore[no-redef]
+        def set_attribute(self, key: str, value: Any) -> None:
+            return None
 
-        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(
+        def record_exception(
             self,
-            arrays: list[Any],
-            names: Optional[list[str]] = None,
-            schema: Optional[Any] = None,
-            metadata: Optional[Mapping[str, Any]] = None,
-        ) -> Any: ...
-        def from_pydict(
+            exception: "Exception",
+            attributes: "Optional[Mapping[str, Any]]" = None,
+            timestamp: "Optional[int]" = None,
+            escaped: bool = False,
+        ) -> None:
+            return None
+
+        def set_status(self, status: Any, description: "Optional[str]" = None) -> None:
+            return None
+
+        def end(self, end_time: "Optional[int]" = None) -> None:
+            return None
+
+        def __enter__(self) -> "Span":
+            return self  # type: ignore[return-value]
+
+        def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+            return None
+
+    class Tracer:  # type: ignore[no-redef]
+        def start_span(
             self,
-            mapping: dict[str, Any],
-            schema: Optional[Any] = None,
-            metadata: Optional[Mapping[str, Any]] = None,
+            name: str,
+            context: Any = None,
+            kind: Any = None,
+            attributes: Any = None,
+            links: Any = None,
+            start_time: Any = None,
+            record_exception: bool = True,
+            set_status_on_exception: bool = True,
+        ) -> Span:
+            return Span()  # type: ignore[abstract]
+
+    class _TraceModule:
+        def get_tracer(
+            self,
+            instrumenting_module_name: str,
+            instrumenting_library_version: "Optional[str]" = None,
+            schema_url: "Optional[str]" = None,
+            tracer_provider: Any = None,
+        ) -> Tracer:
+            return Tracer()  # type: ignore[abstract] # pragma: no cover
+
+        TracerProvider = type(None)  # Shim for TracerProvider if needed elsewhere
+        StatusCode = type(None)  # Shim for StatusCode
+        Status = type(None)  # Shim for Status
+
+    trace = _TraceModule()  # type: ignore[assignment]
+    StatusCode = trace.StatusCode  # type: ignore[misc]
+    Status = trace.Status  # type: ignore[misc]
+    OPENTELEMETRY_INSTALLED = False  # pyright: ignore[reportConstantRedefinition]
+
+
+try:
+    from prometheus_client import (  # pyright: ignore[reportMissingImports, reportAssignmentType]
+        Counter,  # pyright: ignore[reportAssignmentType]
+        Gauge,  # pyright: ignore[reportAssignmentType]
+        Histogram,  # pyright: ignore[reportAssignmentType]
+    )
+
+    PROMETHEUS_INSTALLED = True
+except ImportError:
+    # Define shims for when prometheus_client is not installed
+
+    class _Metric:  # Base shim for metrics
+        def __init__(
+            self,
+            name: str,
+            documentation: str,
+            labelnames: tuple[str, ...] = (),
+            namespace: str = "",
+            subsystem: str = "",
+            unit: str = "",
+            registry: Any = None,
+            ejemplar_fn: Any = None,
+        ) -> None:
+            return None
+
+        def labels(self, *labelvalues: str, **labelkwargs: str) -> "_MetricInstance":
+            return _MetricInstance()
+
+    class _MetricInstance:
+        def inc(self, amount: float = 1) -> None:
+            return None
+
+        def dec(self, amount: float = 1) -> None:
+            return None
+
+        def set(self, value: float) -> None:
+            return None
+
+        def observe(self, amount: float) -> None:
+            return None
+
+    class Counter(_Metric):  # type: ignore[no-redef]
+        def labels(self, *labelvalues: str, **labelkwargs: str) -> _MetricInstance:
+            return _MetricInstance()  # pragma: no cover
+
+    class Gauge(_Metric):  # type: ignore[no-redef]
+        def labels(self, *labelvalues: str, **labelkwargs: str) -> _MetricInstance:
+            return _MetricInstance()  # pragma: no cover
+
+    class Histogram(_Metric):  # type: ignore[no-redef]
+        def labels(self, *labelvalues: str, **labelkwargs: str) -> _MetricInstance:
+            return _MetricInstance()  # pragma: no cover
+
+    PROMETHEUS_INSTALLED = False  # pyright: ignore[reportConstantRedefinition]
+
+
+try:
+    import aiosql  # pyright: ignore[reportMissingImports, reportAssignmentType]
+    from aiosql.types import (  # pyright: ignore[reportMissingImports, reportAssignmentType]
+        AsyncDriverAdapterProtocol as AiosqlAsyncProtocol,  # pyright: ignore[reportMissingImports, reportAssignmentType]
+    )
+    from aiosql.types import (  # pyright: ignore[reportMissingImports, reportAssignmentType]
+        DriverAdapterProtocol as AiosqlProtocol,  # pyright: ignore[reportMissingImports, reportAssignmentType]
+    )
+    from aiosql.types import ParamType as AiosqlParamType  # pyright: ignore[reportMissingImports, reportAssignmentType]
+    from aiosql.types import (
+        SQLOperationType as AiosqlSQLOperationType,  # pyright: ignore[reportMissingImports, reportAssignmentType]
+    )
+    from aiosql.types import (  # pyright: ignore[reportMissingImports, reportAssignmentType]
+        SyncDriverAdapterProtocol as AiosqlSyncProtocol,  # pyright: ignore[reportMissingImports, reportAssignmentType]
+    )
+
+    AIOSQL_INSTALLED = True
+except ImportError:
+    # Define shims for when aiosql is not installed
+
+    class _AiosqlShim:
+        """Placeholder aiosql module"""
+
+        @staticmethod
+        def from_path(sql_path: str, driver_adapter: Any, **kwargs: Any) -> Any:
+            """Placeholder from_path method"""
+            return None  # pragma: no cover
+
+        @staticmethod
+        def from_str(sql_str: str, driver_adapter: Any, **kwargs: Any) -> Any:
+            """Placeholder from_str method"""
+            return None  # pragma: no cover
+
+    aiosql = _AiosqlShim()  # type: ignore[assignment]
+
+    # Placeholder types for aiosql protocols
+    AiosqlParamType = Union[dict[str, Any], list[Any], tuple[Any, ...], None]  # type: ignore[misc]
+
+    class AiosqlSQLOperationType(Enum):  # type: ignore[no-redef]
+        """Enumeration of aiosql operation types."""
+
+        INSERT_RETURNING = 0
+        INSERT_UPDATE_DELETE = 1
+        INSERT_UPDATE_DELETE_MANY = 2
+        SCRIPT = 3
+        SELECT = 4
+        SELECT_ONE = 5
+        SELECT_VALUE = 6
+
+    @runtime_checkable
+    class AiosqlProtocol(Protocol):  # type: ignore[no-redef]
+        """Placeholder for aiosql DriverAdapterProtocol"""
+
+        def process_sql(self, query_name: str, op_type: Any, sql: str) -> str: ...
+
+    @runtime_checkable
+    class AiosqlSyncProtocol(Protocol):  # type: ignore[no-redef]
+        """Placeholder for aiosql SyncDriverAdapterProtocol"""
+
+        is_aio_driver: "ClassVar[bool]"
+
+        def process_sql(self, query_name: str, op_type: Any, sql: str) -> str: ...
+        def select(
+            self, conn: Any, query_name: str, sql: str, parameters: Any, record_class: "Optional[Any]" = None
         ) -> Any: ...
-        def from_batches(self, batches: Iterable[Any], schema: Optional[Any] = None) -> Any: ...
+        def select_one(
+            self, conn: Any, query_name: str, sql: str, parameters: Any, record_class: "Optional[Any]" = None
+        ) -> "Optional[Any]": ...
+        def select_value(self, conn: Any, query_name: str, sql: str, parameters: Any) -> "Optional[Any]": ...
+        def select_cursor(self, conn: Any, query_name: str, sql: str, parameters: Any) -> Any: ...
+        def insert_update_delete(self, conn: Any, query_name: str, sql: str, parameters: Any) -> int: ...
+        def insert_update_delete_many(self, conn: Any, query_name: str, sql: str, parameters: Any) -> int: ...
+        def insert_returning(self, conn: Any, query_name: str, sql: str, parameters: Any) -> "Optional[Any]": ...
 
-    PYARROW_INSTALLED = False  # pyright: ignore[reportConstantRedefinition]
+    @runtime_checkable
+    class AiosqlAsyncProtocol(Protocol):  # type: ignore[no-redef]
+        """Placeholder for aiosql AsyncDriverAdapterProtocol"""
+
+        is_aio_driver: "ClassVar[bool]"
+
+        def process_sql(self, query_name: str, op_type: Any, sql: str) -> str: ...
+        async def select(
+            self, conn: Any, query_name: str, sql: str, parameters: Any, record_class: "Optional[Any]" = None
+        ) -> Any: ...
+        async def select_one(
+            self, conn: Any, query_name: str, sql: str, parameters: Any, record_class: "Optional[Any]" = None
+        ) -> "Optional[Any]": ...
+        async def select_value(self, conn: Any, query_name: str, sql: str, parameters: Any) -> "Optional[Any]": ...
+        async def select_cursor(self, conn: Any, query_name: str, sql: str, parameters: Any) -> Any: ...
+        async def insert_update_delete(self, conn: Any, query_name: str, sql: str, parameters: Any) -> None: ...
+        async def insert_update_delete_many(self, conn: Any, query_name: str, sql: str, parameters: Any) -> None: ...
+        async def insert_returning(self, conn: Any, query_name: str, sql: str, parameters: Any) -> "Optional[Any]": ...
+
+    AIOSQL_INSTALLED = False  # pyright: ignore[reportConstantRedefinition]
+
+
+FSSPEC_INSTALLED = bool(find_spec("fsspec"))
+OBSTORE_INSTALLED = bool(find_spec("obstore"))
+PGVECTOR_INSTALLED = bool(find_spec("pgvector"))
 
 
 __all__ = (
+    "AIOSQL_INSTALLED",
+    "FSSPEC_INSTALLED",
     "LITESTAR_INSTALLED",
     "MSGSPEC_INSTALLED",
+    "OBSTORE_INSTALLED",
+    "OPENTELEMETRY_INSTALLED",
+    "PGVECTOR_INSTALLED",
+    "PROMETHEUS_INSTALLED",
     "PYARROW_INSTALLED",
     "PYDANTIC_INSTALLED",
     "UNSET",
+    "AiosqlAsyncProtocol",
+    "AiosqlParamType",
+    "AiosqlProtocol",
+    "AiosqlSQLOperationType",
+    "AiosqlSyncProtocol",
+    "ArrowRecordBatch",
+    "ArrowRecordBatchResult",
     "ArrowTable",
+    "ArrowTableResult",
     "BaseModel",
+    "Counter",
     "DTOData",
     "DataclassProtocol",
     "Empty",
     "EmptyEnum",
     "EmptyType",
     "FailFast",
+    "Gauge",
+    "Histogram",
+    "Span",
+    "Status",
+    "StatusCode",
     "Struct",
+    "T",
+    "T_co",
+    "Tracer",
     "TypeAdapter",
     "UnsetType",
+    "aiosql",
     "convert",
+    "trace",
 )

sqlspec/adapters/adbc/__init__.py

@@ -1,8 +1,4 @@
-from sqlspec.adapters.adbc.config import AdbcConfig
+from sqlspec.adapters.adbc.config import CONNECTION_FIELDS, AdbcConfig
 from sqlspec.adapters.adbc.driver import AdbcConnection, AdbcDriver
 
-__all__ = (
-    "AdbcConfig",
-    "AdbcConnection",
-    "AdbcDriver",
-)
+__all__ = ("CONNECTION_FIELDS", "AdbcConfig", "AdbcConnection", "AdbcDriver")