sqlspec 0.16.1__cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/storage/registry.py

@@ -0,0 +1,239 @@
+"""Unified Storage Registry for ObjectStore backends.
+
+Provides a flexible, lazy-loading storage registry that supports URI-first access
+pattern with automatic backend detection, ObStore preferred with FSSpec fallback,
+intelligent 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 mypy_extensions import mypyc_attr
+
+from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError
+from sqlspec.protocols import ObjectStoreProtocol
+from sqlspec.storage.capabilities import StorageCapabilities
+from sqlspec.typing import FSSPEC_INSTALLED, OBSTORE_INSTALLED
+
+__all__ = ("StorageRegistry", "storage_registry")
+
+logger = logging.getLogger(__name__)
+
+
+SCHEME_REGEX: Final = re.compile(r"([a-zA-Z0-9+.-]+)://")
+FILE_PROTOCOL: Final[str] = "file"
+S3_PROTOCOL: Final[str] = "s3"
+GCS_PROTOCOL: Final[str] = "gs"
+AZURE_PROTOCOL: Final[str] = "az"
+FSSPEC_ONLY_SCHEMES: Final[frozenset[str]] = frozenset({"http", "https", "ftp", "sftp", "ssh"})
+
+
+@mypyc_attr(allow_interpreted_subclasses=True)
+class StorageRegistry:
+    """Storage registry with URI-first access and automatic backend selection.
+
+    Provides URI-first access pattern with automatic backend selection.
+    Named aliases support complex configurations.
+
+    Examples:
+        backend = registry.get("s3://my-bucket/file.parquet")
+        backend = registry.get("file:///tmp/data.csv")
+        backend = registry.get("gs://bucket/data.json")
+
+        registry.register_alias(
+            "production-s3",
+            uri="s3://prod-bucket/data",
+            base_path="sqlspec",
+            aws_access_key_id="...",
+            aws_secret_access_key="..."
+        )
+        backend = registry.get("production-s3")
+    """
+
+    __slots__ = ("_alias_configs", "_aliases", "_cache", "_instances")
+
+    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._cache: dict[str, tuple[str, type[ObjectStoreProtocol]]] = {}
+
+    def register_alias(
+        self,
+        alias: str,
+        uri: str,
+        *,
+        backend: Optional[type[ObjectStoreProtocol]] = None,
+        base_path: str = "",
+        config: Optional[dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Register a named alias for a storage configuration.
+
+        Args:
+            alias: Unique alias name for the configuration
+            uri: Storage URI (e.g., "s3://bucket", "file:///path")
+            backend: Backend class to use (auto-detected from URI if not provided)
+            base_path: Base path to prepend to all operations
+            config: Additional configuration dict
+            **kwargs: Backend-specific configuration options
+        """
+        if backend is None:
+            backend = self._determine_backend_class(uri)
+
+        config = config or {}
+        config.update(kwargs)
+        backend_config = dict(config)
+        if base_path:
+            backend_config["base_path"] = base_path
+        self._alias_configs[alias] = (backend, uri, backend_config)
+        test_config = dict(backend_config)
+        test_config["uri"] = uri
+        self._aliases[alias] = test_config
+
+    def get(self, uri_or_alias: Union[str, Path], **kwargs: Any) -> ObjectStoreProtocol:
+        """Get backend instance using URI-first routing with automatic backend selection.
+
+        Args:
+            uri_or_alias: URI to resolve directly OR named alias
+            **kwargs: Additional backend-specific configuration options
+
+        Returns:
+            Backend instance with automatic backend selection
+
+        Raises:
+            ImproperConfigurationError: If alias not found or invalid input
+        """
+        if not uri_or_alias:
+            msg = "URI or alias cannot be empty."
+            raise ImproperConfigurationError(msg)
+
+        if isinstance(uri_or_alias, Path):
+            uri_or_alias = f"file://{uri_or_alias.resolve()}"
+
+        cache_key = (uri_or_alias, tuple(sorted(kwargs.items()))) if kwargs 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 (
+            Path(uri_or_alias).exists()
+            or Path(uri_or_alias).is_absolute()
+            or uri_or_alias.startswith(("~", "."))
+            or ":\\" in uri_or_alias
+            or "/" in uri_or_alias
+        ):
+            scheme = "file"
+            uri_or_alias = f"file://{uri_or_alias}"
+
+        if scheme:
+            instance = self._resolve_from_uri(uri_or_alias, **kwargs)
+        elif uri_or_alias in self._alias_configs:
+            backend_cls, stored_uri, config = self._alias_configs[uri_or_alias]
+            instance = backend_cls(stored_uri, **{**config, **kwargs})
+        else:
+            msg = f"Unknown storage alias or invalid URI: '{uri_or_alias}'"
+            raise ImproperConfigurationError(msg)
+        self._instances[cache_key] = instance
+        return instance
+
+    def _resolve_from_uri(self, uri: str, **kwargs: Any) -> ObjectStoreProtocol:
+        """Resolve backend from URI, trying ObStore first, then FSSpec."""
+        scheme = self._get_scheme(uri)
+        if scheme not in FSSPEC_ONLY_SCHEMES and OBSTORE_INSTALLED:
+            try:
+                return self._create_backend("obstore", uri, **kwargs)
+            except (ValueError, ImportError, NotImplementedError):
+                pass
+        if FSSPEC_INSTALLED:
+            try:
+                return self._create_backend("fsspec", uri, **kwargs)
+            except (ValueError, ImportError, NotImplementedError):
+                pass
+        msg = f"No storage backend available for scheme '{scheme}'. Install obstore or fsspec."
+        raise MissingDependencyError(msg)
+
+    def _determine_backend_class(self, uri: str) -> type[ObjectStoreProtocol]:
+        """Determine the backend class for a URI based on availability."""
+        scheme = self._get_scheme(uri)
+        if scheme in FSSPEC_ONLY_SCHEMES and FSSPEC_INSTALLED:
+            return self._get_backend_class("fsspec")
+        if OBSTORE_INSTALLED:
+            return self._get_backend_class("obstore")
+        if FSSPEC_INSTALLED:
+            return self._get_backend_class("fsspec")
+        msg = f"No backend available for URI scheme '{scheme}'. Install obstore or fsspec."
+        raise MissingDependencyError(msg)
+
+    def _get_backend_class(self, backend_type: str) -> type[ObjectStoreProtocol]:
+        """Get backend class by type name."""
+        if backend_type == "obstore":
+            from sqlspec.storage.backends.obstore import ObStoreBackend
+
+            return cast("type[ObjectStoreProtocol]", ObStoreBackend)
+        if backend_type == "fsspec":
+            from sqlspec.storage.backends.fsspec import FSSpecBackend
+
+            return cast("type[ObjectStoreProtocol]", FSSpecBackend)
+        msg = f"Unknown backend type: {backend_type}. Supported types: 'obstore', 'fsspec'"
+        raise ValueError(msg)
+
+    def _create_backend(self, backend_type: str, uri: str, **kwargs: Any) -> ObjectStoreProtocol:
+        """Create backend instance for URI."""
+        return self._get_backend_class(backend_type)(uri, **kwargs)
+
+    def _get_scheme(self, uri: str) -> Optional[str]:
+        """Extract the scheme from a URI using regex."""
+        if not uri:
+            return None
+        match = SCHEME_REGEX.match(uri)
+        return match.group(1).lower() if match else None
+
+    def is_alias_registered(self, alias: str) -> bool:
+        """Check if a named alias is registered."""
+        return alias in self._alias_configs
+
+    def list_aliases(self) -> list[str]:
+        """List all registered aliases."""
+        return list(self._alias_configs.keys())
+
+    def clear_cache(self, uri_or_alias: Optional[str] = None) -> None:
+        """Clear resolved backend cache."""
+        if uri_or_alias:
+            self._instances.pop(uri_or_alias, None)
+        else:
+            self._instances.clear()
+
+    def clear(self) -> None:
+        """Clear all aliases and instances."""
+        self._alias_configs.clear()
+        self._aliases.clear()
+        self._instances.clear()
+
+    def clear_instances(self) -> None:
+        """Clear only cached instances, keeping aliases."""
+        self._instances.clear()
+
+    def clear_aliases(self) -> None:
+        """Clear only aliases, keeping cached instances."""
+        self._alias_configs.clear()
+        self._aliases.clear()
+
+    def get_backend_capabilities(self, uri_or_alias: Union[str, Path]) -> "StorageCapabilities":
+        """Get capabilities for a backend without creating an instance."""
+        if isinstance(uri_or_alias, Path):
+            uri_or_alias = f"file://{uri_or_alias.resolve()}"
+        if "://" in uri_or_alias:
+            backend_cls = self._determine_backend_class(uri_or_alias)
+        elif uri_or_alias in self._alias_configs:
+            backend_cls, _, _ = self._alias_configs[uri_or_alias]
+        else:
+            msg = f"Unknown storage alias or invalid URI: '{uri_or_alias}'"
+            raise ImproperConfigurationError(msg)
+        if hasattr(backend_cls, "capabilities"):
+            return backend_cls.capabilities
+        return StorageCapabilities()
+
+
+storage_registry = StorageRegistry()

sqlspec/typing.py

@@ -0,0 +1,299 @@
+# pyright: ignore[reportAttributeAccessIssue]
+from collections.abc import Iterator, Mapping
+from functools import lru_cache
+from typing import TYPE_CHECKING, Annotated, Any, Protocol, Union
+
+from typing_extensions import TypeAlias, TypeVar
+
+from sqlspec._typing import (
+    AIOSQL_INSTALLED,
+    ATTRS_INSTALLED,
+    CATTRS_INSTALLED,
+    FSSPEC_INSTALLED,
+    LITESTAR_INSTALLED,
+    MSGSPEC_INSTALLED,
+    OBSTORE_INSTALLED,
+    OPENTELEMETRY_INSTALLED,
+    PGVECTOR_INSTALLED,
+    PROMETHEUS_INSTALLED,
+    PYARROW_INSTALLED,
+    PYDANTIC_INSTALLED,
+    DataclassProtocol,
+    Empty,
+    EmptyEnum,
+    EmptyType,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from sqlspec._typing import (
+        UNSET,
+        AiosqlAsyncProtocol,
+        AiosqlParamType,
+        AiosqlProtocol,
+        AiosqlSQLOperationType,
+        AiosqlSyncProtocol,
+        ArrowRecordBatch,
+        ArrowTable,
+        AttrsInstance,
+        AttrsInstanceStub,
+        BaseModel,
+        BaseModelStub,
+        Counter,
+        DTOData,
+        FailFast,
+        Gauge,
+        Histogram,
+        Span,
+        Status,
+        StatusCode,
+        Struct,
+        StructStub,
+        Tracer,
+        TypeAdapter,
+        UnsetType,
+        aiosql,
+        attrs_asdict,
+        attrs_define,
+        attrs_field,
+        attrs_fields,
+        attrs_has,
+        cattrs_structure,
+        cattrs_unstructure,
+        convert,
+        trace,
+    )
+else:
+    from sqlspec._typing import (
+        UNSET,
+        AiosqlAsyncProtocol,
+        AiosqlParamType,
+        AiosqlProtocol,
+        AiosqlSQLOperationType,
+        AiosqlSyncProtocol,
+        ArrowRecordBatch,
+        ArrowTable,
+        AttrsInstance,
+        BaseModel,
+        Counter,
+        DTOData,
+        FailFast,
+        Gauge,
+        Histogram,
+        Span,
+        Status,
+        StatusCode,
+        Struct,
+        Tracer,
+        TypeAdapter,
+        UnsetType,
+        aiosql,
+        attrs_asdict,
+        attrs_define,
+        attrs_field,
+        attrs_fields,
+        attrs_has,
+        cattrs_structure,
+        cattrs_unstructure,
+        convert,
+        trace,
+    )
+
+
+class DictLike(Protocol):
+    """A protocol for objects that behave like a dictionary for reading."""
+
+    def __getitem__(self, key: str) -> Any: ...
+    def __iter__(self) -> Iterator[str]: ...
+    def __len__(self) -> int: ...
+
+
+PYDANTIC_USE_FAILFAST = False
+
+
+if TYPE_CHECKING:
+    T = TypeVar("T")
+    ConnectionT = TypeVar("ConnectionT")
+    """Type variable for connection types.
+
+    :class:`~sqlspec.typing.ConnectionT`
+    """
+    PoolT = TypeVar("PoolT")
+    """Type variable for pool types.
+
+    :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.
+
+    :class:`DictLike` | :class:`msgspec.Struct` | :class:`pydantic.BaseModel` | :class:`DataclassProtocol` | :class:`AttrsInstance`
+    """
+    RowT = TypeVar("RowT", bound="dict[str, Any]")
+else:
+    T = Any
+    ConnectionT = Any
+    PoolT = Any
+    PoolT_co = Any
+    ModelT = Any
+    RowT = 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]"
+"""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]"
+"""Type alias for statement parameters.
+
+Represents:
+- :type:`dict[str, Any]`
+- :type:`list[Any]`
+- :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)
+def get_type_adapter(f: "type[T]") -> Any:
+    """Caches and returns a pydantic type adapter.
+
+    Args:
+        f: Type to create a type adapter for.
+
+    Returns:
+        :class:`pydantic.TypeAdapter`[:class:`typing.TypeVar`[T]]
+    """
+    if PYDANTIC_USE_FAILFAST:
+        return TypeAdapter(Annotated[f, FailFast()])
+    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",
+    "CATTRS_INSTALLED",
+    "FSSPEC_INSTALLED",
+    "LITESTAR_INSTALLED",
+    "MSGSPEC_INSTALLED",
+    "OBSTORE_INSTALLED",
+    "OPENTELEMETRY_INSTALLED",
+    "PGVECTOR_INSTALLED",
+    "PROMETHEUS_INSTALLED",
+    "PYARROW_INSTALLED",
+    "PYDANTIC_INSTALLED",
+    "PYDANTIC_USE_FAILFAST",
+    "UNSET",
+    "AiosqlAsyncProtocol",
+    "AiosqlParamType",
+    "AiosqlProtocol",
+    "AiosqlSQLOperationType",
+    "AiosqlSyncProtocol",
+    "ArrowRecordBatch",
+    "ArrowTable",
+    "AttrsInstance",
+    "BaseModel",
+    "BulkModelDict",
+    "ConnectionT",
+    "Counter",
+    "DTOData",
+    "DataclassProtocol",
+    "DictLike",
+    "DictRow",
+    "Empty",
+    "EmptyEnum",
+    "EmptyType",
+    "FailFast",
+    "Gauge",
+    "Histogram",
+    "Mapping",
+    "MixinOf",
+    "ModelDTOT",
+    "ModelDict",
+    "ModelDict",
+    "ModelDictList",
+    "ModelDictList",
+    "ModelT",
+    "PoolT",
+    "PoolT_co",
+    "PydanticOrMsgspecT",
+    "RowT",
+    "Span",
+    "StatementParameters",
+    "Status",
+    "StatusCode",
+    "Struct",
+    "SupportedSchemaModel",
+    "Tracer",
+    "TupleRow",
+    "TypeAdapter",
+    "UnsetType",
+    "aiosql",
+    "attrs_asdict",
+    "attrs_define",
+    "attrs_field",
+    "attrs_fields",
+    "attrs_has",
+    "cattrs_structure",
+    "cattrs_unstructure",
+    "convert",
+    "get_type_adapter",
+    "trace",
+)

sqlspec/utils/__init__.py

@@ -0,0 +1,3 @@
+from sqlspec.utils import deprecation, fixtures, module_loader, singleton, sync_tools, text, type_guards
+
+__all__ = ("deprecation", "fixtures", "module_loader", "singleton", "sync_tools", "text", "type_guards")

sqlspec/utils/correlation.py

@@ -0,0 +1,150 @@
+"""Correlation ID tracking for distributed tracing.
+
+This module provides utilities for tracking correlation IDs across
+database operations, enabling distributed tracing and debugging.
+"""
+
+from __future__ import annotations
+
+import uuid
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Generator, MutableMapping
+    from logging import LoggerAdapter
+
+__all__ = ("CorrelationContext", "correlation_context", "get_correlation_adapter")
+
+
+class CorrelationContext:
+    """Context manager for correlation ID tracking.
+
+    This class provides a context-aware way to track correlation IDs
+    across async and sync operations.
+    """
+
+    _correlation_id: ContextVar[str | None] = ContextVar("sqlspec_correlation_id", default=None)
+
+    @classmethod
+    def get(cls) -> str | None:
+        """Get the current correlation ID.
+
+        Returns:
+            The current correlation ID or None if not set
+        """
+        return cls._correlation_id.get()
+
+    @classmethod
+    def set(cls, correlation_id: str | None) -> None:
+        """Set the correlation ID.
+
+        Args:
+            correlation_id: The correlation ID to set
+        """
+        cls._correlation_id.set(correlation_id)
+
+    @classmethod
+    def generate(cls) -> str:
+        """Generate a new correlation ID.
+
+        Returns:
+            A new UUID-based correlation ID
+        """
+        return str(uuid.uuid4())
+
+    @classmethod
+    @contextmanager
+    def context(cls, correlation_id: str | None = None) -> Generator[str, None, None]:
+        """Context manager for correlation ID scope.
+
+        Args:
+            correlation_id: The correlation ID to use. If None, generates a new one.
+
+        Yields:
+            The correlation ID being used
+        """
+        if correlation_id is None:
+            correlation_id = cls.generate()
+
+        previous_id = cls.get()
+
+        try:
+            cls.set(correlation_id)
+            yield correlation_id
+        finally:
+            cls.set(previous_id)
+
+    @classmethod
+    def clear(cls) -> None:
+        """Clear the current correlation ID."""
+        cls.set(None)
+
+    @classmethod
+    def to_dict(cls) -> dict[str, Any]:
+        """Get correlation context as a dictionary.
+
+        Returns:
+            Dictionary with correlation_id key if set
+        """
+        correlation_id = cls.get()
+        return {"correlation_id": correlation_id} if correlation_id else {}
+
+
+@contextmanager
+def correlation_context(correlation_id: str | None = None) -> Generator[str, None, None]:
+    """Convenience context manager for correlation ID tracking.
+
+    Args:
+        correlation_id: Optional correlation ID. If None, generates a new one.
+
+    Yields:
+        The active correlation ID
+
+    Example:
+        ```python
+        with correlation_context() as correlation_id:
+            logger.info(
+                "Processing request",
+                extra={"correlation_id": correlation_id},
+            )
+        ```
+    """
+    with CorrelationContext.context(correlation_id) as cid:
+        yield cid
+
+
+def get_correlation_adapter(logger: Any) -> LoggerAdapter:
+    """Get a logger adapter that automatically includes correlation ID.
+
+    Args:
+        logger: The base logger to wrap
+
+    Returns:
+        LoggerAdapter that includes correlation ID in all logs
+    """
+    from logging import 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]]:
+            """Add correlation ID to the log record.
+
+            Args:
+                msg: The log message
+                kwargs: Keyword arguments for the log record
+
+            Returns:
+                The message and updated kwargs
+            """
+            extra = kwargs.get("extra", {})
+
+            if correlation_id := CorrelationContext.get():
+                extra["correlation_id"] = correlation_id
+
+            kwargs["extra"] = extra
+            return msg, dict(kwargs)
+
+    return CorrelationAdapter(logger, {})