sqlspec 0.26.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

sqlspec/storage/registry.py

@@ -8,33 +8,19 @@ scheme-based routing, and named aliases for common configurations.
 import logging
 import re
 from pathlib import Path
-from typing import Any, Final, Optional, Union, cast
+from typing import Any, Final, cast
 
 from mypy_extensions import mypyc_attr
 
 from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError
 from sqlspec.protocols import ObjectStoreProtocol
 from sqlspec.typing import FSSPEC_INSTALLED, OBSTORE_INSTALLED
+from sqlspec.utils.type_guards import is_local_path
 
 __all__ = ("StorageRegistry", "storage_registry")
 
 logger = logging.getLogger(__name__)
 
-
-def _is_local_uri(uri: str) -> bool:
-    """Check if URI represents a local filesystem path."""
-    if "://" in uri and not uri.startswith("file://"):
-        return False
-    windows_drive_min_length = 3
-    return (
-        Path(uri).exists()
-        or Path(uri).is_absolute()
-        or uri.startswith(("~", ".", "/"))
-        or (len(uri) >= windows_drive_min_length and uri[1:3] == ":\\")
-        or "/" in uri
-    )
-
-
 SCHEME_REGEX: Final = re.compile(r"([a-zA-Z0-9+.-]+)://")
 
 
@@ -74,7 +60,7 @@ class StorageRegistry:
     def __init__(self) -> None:
         self._alias_configs: dict[str, tuple[type[ObjectStoreProtocol], str, dict[str, Any]]] = {}
         self._aliases: dict[str, dict[str, Any]] = {}
-        self._instances: dict[Union[str, tuple[str, tuple[tuple[str, Any], ...]]], ObjectStoreProtocol] = {}
+        self._instances: dict[str | tuple[str, tuple[tuple[str, Any], ...]], ObjectStoreProtocol] = {}
         self._cache: dict[str, tuple[str, type[ObjectStoreProtocol]]] = {}
 
     def _make_hashable(self, obj: Any) -> Any:
@@ -88,7 +74,7 @@ class StorageRegistry:
         return obj
 
     def register_alias(
-        self, alias: str, uri: str, *, backend: Optional[str] = None, base_path: str = "", **kwargs: Any
+        self, alias: str, uri: str, *, backend: str | None = None, base_path: str = "", **kwargs: Any
     ) -> None:
         """Register a named alias for a storage configuration.
 
@@ -110,9 +96,7 @@ class StorageRegistry:
         test_config["uri"] = uri
         self._aliases[alias] = test_config
 
-    def get(
-        self, uri_or_alias: Union[str, Path], *, backend: Optional[str] = None, **kwargs: Any
-    ) -> ObjectStoreProtocol:
+    def get(self, uri_or_alias: str | Path, *, backend: str | None = None, **kwargs: Any) -> ObjectStoreProtocol:
         """Get backend instance using URI-first routing with automatic backend selection.
 
         Args:
@@ -133,11 +117,15 @@ class StorageRegistry:
         if isinstance(uri_or_alias, Path):
             uri_or_alias = f"file://{uri_or_alias.resolve()}"
 
-        cache_key = (uri_or_alias, self._make_hashable(kwargs)) if kwargs else uri_or_alias
+        # Include backend in cache key to ensure different backends for same URI are cached separately
+        cache_params = dict(kwargs)
+        if backend:
+            cache_params["__backend__"] = backend
+        cache_key = (uri_or_alias, self._make_hashable(cache_params)) if cache_params else uri_or_alias
         if cache_key in self._instances:
             return self._instances[cache_key]
         scheme = self._get_scheme(uri_or_alias)
-        if not scheme and _is_local_uri(uri_or_alias):
+        if not scheme and is_local_path(uri_or_alias):
             scheme = "file"
             uri_or_alias = f"file://{uri_or_alias}"
 
@@ -154,57 +142,105 @@ class StorageRegistry:
         self._instances[cache_key] = instance
         return instance
 
-    def _resolve_from_uri(
-        self, uri: str, *, backend_override: Optional[str] = None, **kwargs: Any
-    ) -> ObjectStoreProtocol:
-        """Resolve backend from URI with optional backend override."""
+    def _resolve_from_uri(self, uri: str, *, backend_override: str | None = None, **kwargs: Any) -> ObjectStoreProtocol:
+        """Resolve backend from URI with optional backend override.
+
+        Backend selection priority for local files (file:// or bare paths):
+        1. obstore (if installed) - provides async I/O performance
+        2. fsspec (if installed) - async wrapper fallback
+        3. local (always available) - zero-dependency sync backend
+
+        For cloud storage, prefer obstore over fsspec when available.
+
+        Args:
+            uri: Storage URI to resolve.
+            backend_override: Force specific backend type.
+            **kwargs: Additional backend configuration.
+
+        Returns:
+            Configured backend instance.
+
+        Raises:
+            MissingDependencyError: No backend available for URI scheme.
+        """
         if backend_override:
             return self._create_backend(backend_override, uri, **kwargs)
+
         scheme = self._get_scheme(uri)
 
-        # For local files, prefer LocalStore first
+        # NEW: Prefer obstore for local files when available
         if scheme in {None, "file"}:
+            # Try obstore first for async performance
+            if OBSTORE_INSTALLED:
+                try:
+                    return self._create_backend("obstore", uri, **kwargs)
+                except (ValueError, ImportError, NotImplementedError):
+                    pass
+
+            # Fallback to fsspec if available
+            if FSSPEC_INSTALLED:
+                try:
+                    return self._create_backend("fsspec", uri, **kwargs)
+                except (ValueError, ImportError, NotImplementedError):
+                    pass
+
+            # Final fallback: local zero-dependency backend
             return self._create_backend("local", uri, **kwargs)
 
-        # Try ObStore first if available and appropriate
+        # For cloud schemes, prefer obstore over fsspec
         if scheme not in FSSPEC_ONLY_SCHEMES and OBSTORE_INSTALLED:
             try:
                 return self._create_backend("obstore", uri, **kwargs)
             except (ValueError, ImportError, NotImplementedError):
                 pass
 
-        # Try FSSpec if available
+        # Try fsspec if available
         if FSSPEC_INSTALLED:
             try:
                 return self._create_backend("fsspec", uri, **kwargs)
             except (ValueError, ImportError, NotImplementedError):
                 pass
 
-        # For cloud schemes without backends, provide helpful error
+        # No backend available
         msg = f"No backend available for URI scheme '{scheme}'. Install obstore or fsspec for cloud storage support."
         raise MissingDependencyError(msg)
 
     def _determine_backend_class(self, uri: str) -> type[ObjectStoreProtocol]:
-        """Determine the backend class for a URI based on availability."""
+        """Determine the backend class for a URI based on availability.
+
+        Args:
+            uri: Storage URI to analyze.
+
+        Returns:
+            Backend class type to use.
+
+        Raises:
+            MissingDependencyError: No backend available for URI scheme.
+        """
         scheme = self._get_scheme(uri)
 
-        # For local files, always use LocalStore
+        # NEW: For local files, prefer obstore > fsspec > local
         if scheme in {None, "file"}:
+            if OBSTORE_INSTALLED:
+                return self._get_backend_class("obstore")
+            if FSSPEC_INSTALLED:
+                return self._get_backend_class("fsspec")
             return self._get_backend_class("local")
 
         # FSSpec-only schemes require FSSpec
-        if scheme in FSSPEC_ONLY_SCHEMES and FSSPEC_INSTALLED:
+        if scheme in FSSPEC_ONLY_SCHEMES:
+            if not FSSPEC_INSTALLED:
+                msg = f"Scheme '{scheme}' requires fsspec. Install with: pip install fsspec"
+                raise MissingDependencyError(msg)
             return self._get_backend_class("fsspec")
 
-        # Prefer ObStore for cloud storage if available
+        # For cloud schemes, prefer obstore
         if OBSTORE_INSTALLED:
             return self._get_backend_class("obstore")
 
-        # Fall back to FSSpec if available
         if FSSPEC_INSTALLED:
             return self._get_backend_class("fsspec")
 
-        # For cloud schemes without backends, provide helpful error
         msg = f"No backend available for URI scheme '{scheme}'. Install obstore or fsspec for cloud storage support."
         raise MissingDependencyError(msg)
 
@@ -229,7 +265,7 @@ class StorageRegistry:
         """Create backend instance for URI."""
         return self._get_backend_class(backend_type)(uri, **kwargs)
 
-    def _get_scheme(self, uri: str) -> Optional[str]:
+    def _get_scheme(self, uri: str) -> str | None:
         """Extract the scheme from a URI using regex."""
         if not uri:
             return None
@@ -244,7 +280,7 @@ class StorageRegistry:
         """List all registered aliases."""
         return list(self._alias_configs.keys())
 
-    def clear_cache(self, uri_or_alias: Optional[str] = None) -> None:
+    def clear_cache(self, uri_or_alias: str | None = None) -> None:
         """Clear resolved backend cache."""
         if uri_or_alias:
             self._instances.pop(uri_or_alias, None)

sqlspec/typing.py

@@ -1,9 +1,9 @@
 # pyright: ignore[reportAttributeAccessIssue]
-from collections.abc import Iterator, Mapping
+from collections.abc import Iterator
 from functools import lru_cache
-from typing import TYPE_CHECKING, Annotated, Any, Protocol, Union
+from typing import Annotated, Any, Protocol, TypeAlias, _TypedDict  # pyright: ignore
 
-from typing_extensions import TypeAlias, TypeVar
+from typing_extensions import TypeVar
 
 from sqlspec._typing import (
     AIOSQL_INSTALLED,
@@ -16,17 +16,22 @@ from sqlspec._typing import (
     OBSTORE_INSTALLED,
     OPENTELEMETRY_INSTALLED,
     ORJSON_INSTALLED,
+    PANDAS_INSTALLED,
     PGVECTOR_INSTALLED,
+    POLARS_INSTALLED,
     PROMETHEUS_INSTALLED,
     PYARROW_INSTALLED,
     PYDANTIC_INSTALLED,
     UNSET,
     AiosqlAsyncProtocol,
     AiosqlParamType,
-    AiosqlProtocol,
     AiosqlSQLOperationType,
     AiosqlSyncProtocol,
     ArrowRecordBatch,
+    ArrowRecordBatchReader,
+    ArrowRecordBatchReaderProtocol,
+    ArrowSchema,
+    ArrowSchemaProtocol,
     ArrowTable,
     AttrsInstance,
     AttrsInstanceStub,
@@ -41,6 +46,9 @@ from sqlspec._typing import (
     FailFast,
     Gauge,
     Histogram,
+    NumpyArray,
+    PandasDataFrame,
+    PolarsDataFrame,
     Span,
     Status,
     StatusCode,
@@ -61,9 +69,6 @@ from sqlspec._typing import (
     trace,
 )
 
-if TYPE_CHECKING:
-    from collections.abc import Sequence
-
 
 class DictLike(Protocol):
     """A protocol for objects that behave like a dictionary for reading."""
@@ -87,30 +92,22 @@ PoolT = TypeVar("PoolT")
 
 :class:`~sqlspec.typing.PoolT`
 """
-PoolT_co = TypeVar("PoolT_co", covariant=True)
-"""Type variable for covariant pool types.
-
-:class:`~sqlspec.typing.PoolT_co`
-"""
-ModelT = TypeVar("ModelT", bound="Union[DictLike, StructStub, BaseModelStub, DataclassProtocol, AttrsInstanceStub]")
-"""Type variable for model types.
+SchemaT = TypeVar("SchemaT", default=dict[str, Any])
+"""Type variable for schema types (models, TypedDict, dataclasses, etc.).
 
-:class:`DictLike` | :class:`msgspec.Struct` | :class:`pydantic.BaseModel` | :class:`DataclassProtocol` | :class:`AttrsInstance`
+Unbounded TypeVar for use with schema_type parameter in driver methods.
+Supports all schema types including TypedDict which cannot be bounded to a class hierarchy.
 """
-RowT = TypeVar("RowT", bound="dict[str, Any]")
-
 
-DictRow: TypeAlias = "dict[str, Any]"
-"""Type variable for DictRow types."""
-TupleRow: TypeAlias = "tuple[Any, ...]"
-"""Type variable for TupleRow types."""
 
-SupportedSchemaModel: TypeAlias = "Union[DictLike, StructStub, BaseModelStub, DataclassProtocol, AttrsInstanceStub]"
+SupportedSchemaModel: TypeAlias = (
+    DictLike | StructStub | BaseModelStub | DataclassProtocol | AttrsInstanceStub | _TypedDict
+)
 """Type alias for pydantic or msgspec models.
 
 :class:`msgspec.Struct` | :class:`pydantic.BaseModel` | :class:`DataclassProtocol` | :class:`AttrsInstance`
 """
-StatementParameters: TypeAlias = "Union[Any, dict[str, Any], list[Any], tuple[Any, ...], None]"
+StatementParameters: TypeAlias = "Any | dict[str, Any] | list[Any] | tuple[Any, ...] | None"
 """Type alias for statement parameters.
 
 Represents:
@@ -119,40 +116,6 @@ Represents:
 - :type:`tuple[Any, ...]`
 - :type:`None`
 """
-ModelDTOT = TypeVar("ModelDTOT", bound="SupportedSchemaModel")
-"""Type variable for model DTOs.
-
-:class:`msgspec.Struct`|:class:`pydantic.BaseModel`
-"""
-PydanticOrMsgspecT = SupportedSchemaModel
-"""Type alias for pydantic or msgspec models.
-
-:class:`msgspec.Struct` or :class:`pydantic.BaseModel`
-"""
-ModelDict: TypeAlias = (
-    "Union[dict[str, Any], Union[DictLike, StructStub, BaseModelStub, DataclassProtocol, AttrsInstanceStub], Any]"
-)
-"""Type alias for model dictionaries.
-
-Represents:
-- :type:`dict[str, Any]` | :class:`DataclassProtocol` | :class:`msgspec.Struct` |  :class:`pydantic.BaseModel`
-"""
-ModelDictList: TypeAlias = (
-    "Sequence[Union[dict[str, Any], Union[DictLike, StructStub, BaseModelStub, DataclassProtocol, AttrsInstanceStub]]]"
-)
-"""Type alias for model dictionary lists.
-
-A list or sequence of any of the following:
-- :type:`Sequence`[:type:`dict[str, Any]` | :class:`DataclassProtocol` | :class:`msgspec.Struct` | :class:`pydantic.BaseModel`]
-
-"""
-BulkModelDict: TypeAlias = "Union[Sequence[Union[dict[str, Any], Union[DictLike, StructStub, BaseModelStub, DataclassProtocol, AttrsInstanceStub]]], Any]"
-"""Type alias for bulk model dictionaries.
-
-Represents:
-- :type:`Sequence`[:type:`dict[str, Any]` | :class:`DataclassProtocol` | :class:`msgspec.Struct` | :class:`pydantic.BaseModel`]
-- :class:`DTOData`[:type:`list[ModelT]`]
-"""
 
 
 @lru_cache(typed=True)
@@ -170,18 +133,6 @@ def get_type_adapter(f: "type[T]") -> Any:
     return TypeAdapter(f)
 
 
-def MixinOf(base: type[T]) -> type[T]:  # noqa: N802
-    """Useful function to make mixins with baseclass type hint
-
-    ```
-    class StorageMixin(MixinOf(DriverProtocol)): ...
-    ```
-    """
-    if TYPE_CHECKING:
-        return base
-    return type("<MixinOf>", (base,), {})
-
-
 __all__ = (
     "AIOSQL_INSTALLED",
     "ATTRS_INSTALLED",
@@ -193,7 +144,9 @@ __all__ = (
     "OBSTORE_INSTALLED",
     "OPENTELEMETRY_INSTALLED",
     "ORJSON_INSTALLED",
+    "PANDAS_INSTALLED",
     "PGVECTOR_INSTALLED",
+    "POLARS_INSTALLED",
     "PROMETHEUS_INSTALLED",
     "PYARROW_INSTALLED",
     "PYDANTIC_INSTALLED",
@@ -201,38 +154,32 @@ __all__ = (
     "UNSET",
     "AiosqlAsyncProtocol",
     "AiosqlParamType",
-    "AiosqlProtocol",
     "AiosqlSQLOperationType",
     "AiosqlSyncProtocol",
     "ArrowRecordBatch",
+    "ArrowRecordBatchReader",
+    "ArrowRecordBatchReaderProtocol",
+    "ArrowSchema",
+    "ArrowSchemaProtocol",
     "ArrowTable",
     "AttrsInstance",
     "BaseModel",
-    "BulkModelDict",
     "ConnectionT",
     "Counter",
     "DTOData",
     "DataclassProtocol",
     "DictLike",
-    "DictRow",
     "Empty",
     "EmptyEnum",
     "EmptyType",
     "FailFast",
     "Gauge",
     "Histogram",
-    "Mapping",
-    "MixinOf",
-    "ModelDTOT",
-    "ModelDict",
-    "ModelDict",
-    "ModelDictList",
-    "ModelDictList",
-    "ModelT",
+    "NumpyArray",
+    "PandasDataFrame",
+    "PolarsDataFrame",
     "PoolT",
-    "PoolT_co",
-    "PydanticOrMsgspecT",
-    "RowT",
+    "SchemaT",
     "Span",
     "StatementParameters",
     "Status",
@@ -240,7 +187,6 @@ __all__ = (
     "Struct",
     "SupportedSchemaModel",
     "Tracer",
-    "TupleRow",
     "TypeAdapter",
     "UnsetType",
     "aiosql",

sqlspec/utils/__init__.py

@@ -1,10 +1,31 @@
 """Utility functions and classes for SQLSpec.
 
 This package provides various utility modules for deprecation handling,
-fixture loading, logging, module loading, singleton patterns, sync/async
-conversion, text processing, and type guards.
+fixture loading, logging, module loading (including dependency checking),
+portal pattern for async bridging, singleton patterns, sync/async conversion,
+text processing, and type guards.
 """
 
-from sqlspec.utils import deprecation, fixtures, logging, module_loader, singleton, sync_tools, text, type_guards
+from sqlspec.utils import (
+    deprecation,
+    fixtures,
+    logging,
+    module_loader,
+    portal,
+    singleton,
+    sync_tools,
+    text,
+    type_guards,
+)
 
-__all__ = ("deprecation", "fixtures", "logging", "module_loader", "singleton", "sync_tools", "text", "type_guards")
+__all__ = (
+    "deprecation",
+    "fixtures",
+    "logging",
+    "module_loader",
+    "portal",
+    "singleton",
+    "sync_tools",
+    "text",
+    "type_guards",
+)

sqlspec/utils/arrow_helpers.py

@@ -0,0 +1,81 @@
+"""Arrow conversion helpers for dict-to-Arrow transformations.
+
+This module provides utilities for converting Python dictionaries to Apache Arrow
+format, handling empty results, NULL values, and type inference.
+"""
+
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+    from sqlspec.typing import ArrowRecordBatch, ArrowTable
+
+__all__ = ("convert_dict_to_arrow",)
+
+
+def convert_dict_to_arrow(
+    data: "list[dict[str, Any]]",
+    return_format: Literal["table", "reader", "batches"] = "table",
+    batch_size: int | None = None,
+) -> "ArrowTable | ArrowRecordBatch":
+    """Convert list of dictionaries to Arrow Table or RecordBatch.
+
+    Handles empty results, NULL values, and automatic type inference.
+    Used by adapters that don't have native Arrow support to convert
+    dict-based results to Arrow format.
+
+    Args:
+        data: List of dictionaries (one per row).
+        return_format: Output format - "table" for Table, "batches" for RecordBatch.
+            "reader" is converted to "table" (streaming handled at driver level).
+        batch_size: Chunk size for batching (used when return_format="batches").
+
+    Returns:
+        ArrowTable or ArrowRecordBatch depending on return_format.
+
+
+    Examples:
+        >>> data = [
+        ...     {"id": 1, "name": "Alice"},
+        ...     {"id": 2, "name": "Bob"},
+        ... ]
+        >>> table = convert_dict_to_arrow(data, return_format="table")
+        >>> print(table.num_rows)
+        2
+
+        >>> batch = convert_dict_to_arrow(data, return_format="batch")
+        >>> print(batch.num_rows)
+        2
+    """
+    from sqlspec.utils.module_loader import ensure_pyarrow
+
+    ensure_pyarrow()
+
+    import pyarrow as pa
+
+    # Handle empty results
+    if not data:
+        empty_schema = pa.schema([])
+        empty_table = pa.Table.from_pydict({}, schema=empty_schema)
+
+        if return_format == "batches":
+            # Create empty RecordBatch
+            batches = empty_table.to_batches()
+            return batches[0] if batches else pa.RecordBatch.from_pydict({})
+
+        return empty_table
+
+    # Convert list of dicts to columnar format
+    # This is more efficient than row-by-row conversion
+    columns: dict[str, list[Any]] = {key: [row.get(key) for row in data] for key in data[0]}
+
+    # Create Arrow Table (auto type inference)
+    arrow_table = pa.Table.from_pydict(columns)
+
+    # Return appropriate format
+    if return_format == "batches":
+        # Convert table to single RecordBatch
+        batches = arrow_table.to_batches()
+        return batches[0] if batches else pa.RecordBatch.from_pydict({})
+
+    # return_format == "table" or "reader" (reader handled at driver level)
+    return arrow_table

sqlspec/utils/config_resolver.py

@@ -7,7 +7,7 @@ Supports both synchronous and asynchronous callable functions.
 
 import inspect
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Union, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from sqlspec.exceptions import ConfigResolverError
 from sqlspec.utils.module_loader import import_string
@@ -21,7 +21,7 @@ __all__ = ("resolve_config_async", "resolve_config_sync")
 
 async def resolve_config_async(
     config_path: str,
-) -> "Union[list[Union[AsyncDatabaseConfig, SyncDatabaseConfig]], Union[AsyncDatabaseConfig, SyncDatabaseConfig]]":
+) -> "list[AsyncDatabaseConfig | SyncDatabaseConfig] | AsyncDatabaseConfig | SyncDatabaseConfig":
     """Resolve config from dotted path, handling callables and direct instances.
 
     This is the async-first version that handles both sync and async callables efficiently.
@@ -58,7 +58,7 @@ async def resolve_config_async(
 
 def resolve_config_sync(
     config_path: str,
-) -> "Union[list[Union[AsyncDatabaseConfig, SyncDatabaseConfig]], Union[AsyncDatabaseConfig, SyncDatabaseConfig]]":
+) -> "list[AsyncDatabaseConfig | SyncDatabaseConfig] | AsyncDatabaseConfig | SyncDatabaseConfig":
     """Synchronous wrapper for resolve_config.
 
     Args:
@@ -90,7 +90,7 @@ def resolve_config_sync(
 
 def _validate_config_result(
     config_result: Any, config_path: str
-) -> "Union[list[Union[AsyncDatabaseConfig, SyncDatabaseConfig]], Union[AsyncDatabaseConfig, SyncDatabaseConfig]]":
+) -> "list[AsyncDatabaseConfig | SyncDatabaseConfig] | AsyncDatabaseConfig | SyncDatabaseConfig":
     """Validate that the config result is a valid config or list of configs.
 
     Args:
@@ -117,13 +117,13 @@ def _validate_config_result(
                 msg = f"Config '{config_path}' returned invalid config at index {i}. Expected database config instance."
                 raise ConfigResolverError(msg)
 
-        return cast("list[Union[AsyncDatabaseConfig, SyncDatabaseConfig]]", list(config_result))
+        return cast("list[AsyncDatabaseConfig | SyncDatabaseConfig]", list(config_result))
 
     if not _is_valid_config(config_result):
         msg = f"Config '{config_path}' returned invalid type '{type(config_result).__name__}'. Expected database config instance or list."
         raise ConfigResolverError(msg)
 
-    return cast("Union[AsyncDatabaseConfig, SyncDatabaseConfig]", config_result)
+    return cast("AsyncDatabaseConfig | SyncDatabaseConfig", config_result)
 
 
 def _is_valid_config(config: Any) -> bool:

sqlspec/utils/correlation.py

@@ -4,15 +4,14 @@ This module provides utilities for tracking correlation IDs across
 database operations, enabling distributed tracing and debugging.
 """
 
-from __future__ import annotations
-
 import uuid
+from collections.abc import Generator
 from contextlib import contextmanager
 from contextvars import ContextVar
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
-    from collections.abc import Generator, MutableMapping
+    from collections.abc import MutableMapping
     from logging import LoggerAdapter
 
 __all__ = ("CorrelationContext", "correlation_context", "get_correlation_adapter")
@@ -115,7 +114,7 @@ def correlation_context(correlation_id: str | None = None) -> Generator[str, Non
         yield cid
 
 
-def get_correlation_adapter(logger: Any) -> LoggerAdapter:
+def get_correlation_adapter(logger: Any) -> "LoggerAdapter":
     """Get a logger adapter that automatically includes correlation ID.
 
     Args:
@@ -129,7 +128,7 @@ def get_correlation_adapter(logger: Any) -> LoggerAdapter:
     class CorrelationAdapter(LoggerAdapter):
         """Logger adapter that adds correlation ID to all logs."""
 
-        def process(self, msg: str, kwargs: MutableMapping[str, Any]) -> tuple[str, dict[str, Any]]:
+        def process(self, msg: str, kwargs: "MutableMapping[str, Any]") -> tuple[str, dict[str, Any]]:
             """Add correlation ID to the log record.
 
             Args:

sqlspec/utils/data_transformation.py

@@ -5,7 +5,8 @@ field name conversion when mapping database results to schema objects.
 Used primarily for msgspec field name conversion with rename configurations.
 """
 
-from typing import Any, Callable, Union
+from collections.abc import Callable
+from typing import Any
 
 __all__ = ("transform_dict_keys",)
 
@@ -30,7 +31,7 @@ def _safe_convert_key(key: Any, converter: Callable[[str], str]) -> Any:
         return key
 
 
-def transform_dict_keys(data: Union[dict, list, Any], converter: Callable[[str], str]) -> Union[dict, list, Any]:
+def transform_dict_keys(data: dict | list | Any, converter: Callable[[str], str]) -> dict | list | Any:
     """Transform dictionary keys using the provided converter function.
 
     Recursively transforms all dictionary keys in a data structure using