sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/storage/registry.py

@@ -0,0 +1,301 @@
+"""Storage registry for ObjectStore backends.
+
+Provides a storage registry that supports URI-first access
+pattern with automatic backend detection, ObStore preferred with FSSpec fallback,
+scheme-based routing, and named aliases for common configurations.
+"""
+
+import logging
+import re
+from pathlib import Path
+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.logging import get_logger, log_with_context
+from sqlspec.utils.type_guards import is_local_path
+
+__all__ = ("StorageRegistry", "storage_registry")
+
+logger = get_logger(__name__)
+
+SCHEME_REGEX: Final = re.compile(r"([a-zA-Z0-9+.-]+)://")
+
+
+FSSPEC_ONLY_SCHEMES: Final[frozenset[str]] = frozenset({"http", "https", "ftp", "sftp", "ssh"})
+
+
+@mypyc_attr(allow_interpreted_subclasses=True)
+class StorageRegistry:
+    """Global storage registry for named backend configurations.
+
+    Allows registering named storage backends that can be accessed from anywhere
+    in your application. Backends are automatically selected based on URI scheme
+    unless explicitly overridden.
+
+    Examples:
+        backend = registry.get("s3://my-bucket")
+        backend = registry.get("file:///tmp/data")
+        backend = registry.get("gs://my-gcs-bucket")
+
+        registry.register_alias("my_app_store", "file:///tmp/dev_data")
+
+        registry.register_alias("my_app_store", "s3://prod-bucket/data")
+
+        store = registry.get("my_app_store")
+
+        backend = registry.get("s3://bucket", backend="fsspec")
+    """
+
+    __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[str | tuple[str, tuple[tuple[str, Any], ...]], ObjectStoreProtocol] = {}
+        self._cache: dict[str, tuple[str, type[ObjectStoreProtocol]]] = {}
+
+    def _make_hashable(self, obj: Any) -> Any:
+        """Convert nested dict/list structures to hashable tuples."""
+        if isinstance(obj, dict):
+            return tuple(sorted((k, self._make_hashable(v)) for k, v in obj.items()))
+        if isinstance(obj, list):
+            return tuple(self._make_hashable(item) for item in obj)
+        if isinstance(obj, set):
+            return tuple(sorted(self._make_hashable(item) for item in obj))
+        return obj
+
+    def register_alias(
+        self, alias: str, uri: str, *, backend: str | None = None, base_path: str = "", **kwargs: Any
+    ) -> None:
+        """Register a named alias for a storage configuration.
+
+        Args:
+            alias: Unique alias name (e.g., "my_app_store", "user_uploads")
+            uri: Storage URI (e.g., "s3://bucket", "file:///path", "gs://bucket")
+            backend: Force specific backend ("local", "fsspec", "obstore") instead of auto-detection
+            base_path: Base path to prepend to all operations
+            **kwargs: Backend-specific configuration options
+        """
+        backend_cls = self._get_backend_class(backend) if backend else self._determine_backend_class(uri)
+
+        backend_config = dict(kwargs)
+        if base_path:
+            backend_config["base_path"] = base_path
+        self._alias_configs[alias] = (backend_cls, uri, backend_config)
+
+        test_config = dict(backend_config)
+        test_config["uri"] = uri
+        self._aliases[alias] = test_config
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "storage.alias.register",
+            alias=alias,
+            uri=uri,
+            backend_type=backend_cls.__name__,
+            base_path=base_path or None,
+        )
+
+    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:
+            uri_or_alias: URI to resolve directly OR named alias (e.g., "my_app_store")
+            backend: Force specific backend ("local", "fsspec", "obstore") instead of auto-selection
+            **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_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:
+            log_with_context(logger, logging.DEBUG, "storage.resolve", uri_or_alias=str(uri_or_alias), cached=True)
+            return self._instances[cache_key]
+        scheme = self._get_scheme(uri_or_alias)
+        if not scheme and is_local_path(uri_or_alias):
+            scheme = "file"
+            uri_or_alias = f"file://{uri_or_alias}"
+
+        if scheme:
+            instance = self._resolve_from_uri(uri_or_alias, backend_override=backend, **kwargs)
+        elif uri_or_alias in self._alias_configs:
+            backend_cls, stored_uri, config = self._alias_configs[uri_or_alias]
+            if backend:
+                backend_cls = self._get_backend_class(backend)
+            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
+        log_with_context(logger, logging.DEBUG, "storage.resolve", uri_or_alias=str(uri_or_alias), cached=False)
+        return instance
+
+    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)
+
+        if scheme in {None, "file"}:
+            if 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
+
+            return self._create_backend("local", uri, **kwargs)
+
+        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 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.
+
+        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)
+
+        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")
+
+        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")
+
+        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 for cloud storage support."
+        raise MissingDependencyError(msg)
+
+    def _get_backend_class(self, backend_type: str) -> type[ObjectStoreProtocol]:
+        """Get backend class by type name."""
+        if backend_type == "local":
+            from sqlspec.storage.backends.local import LocalStore
+
+            return cast("type[ObjectStoreProtocol]", LocalStore)
+        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: 'local', '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) -> str | None:
+        """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: str | None = 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()
+
+
+storage_registry = StorageRegistry()

sqlspec/typing.py

@@ -0,0 +1,395 @@
+# pyright: ignore[reportAttributeAccessIssue]
+from collections.abc import Iterator
+from functools import lru_cache
+from typing import Annotated, Any, Literal, Protocol, TypeAlias, TypedDict, _TypedDict  # pyright: ignore
+
+from typing_extensions import TypeVar
+
+from sqlspec._typing import (
+    ALLOYDB_CONNECTOR_INSTALLED,
+    ATTRS_INSTALLED,
+    CATTRS_INSTALLED,
+    CLOUD_SQL_CONNECTOR_INSTALLED,
+    FSSPEC_INSTALLED,
+    LITESTAR_INSTALLED,
+    MSGSPEC_INSTALLED,
+    NANOID_INSTALLED,
+    NUMPY_INSTALLED,
+    OBSTORE_INSTALLED,
+    OPENTELEMETRY_INSTALLED,
+    ORJSON_INSTALLED,
+    PANDAS_INSTALLED,
+    PGVECTOR_INSTALLED,
+    POLARS_INSTALLED,
+    PROMETHEUS_INSTALLED,
+    PYARROW_INSTALLED,
+    PYDANTIC_INSTALLED,
+    UNSET,
+    UUID_UTILS_INSTALLED,
+    ArrowRecordBatch,
+    ArrowRecordBatchReader,
+    ArrowRecordBatchReaderProtocol,
+    ArrowSchema,
+    ArrowSchemaProtocol,
+    ArrowTable,
+    AttrsInstance,
+    AttrsInstanceStub,
+    BaseModel,
+    BaseModelStub,
+    Counter,
+    DataclassProtocol,
+    DTOData,
+    Empty,
+    EmptyEnum,
+    EmptyType,
+    FailFast,
+    Gauge,
+    Histogram,
+    NumpyArray,
+    PandasDataFrame,
+    PolarsDataFrame,
+    Span,
+    Status,
+    StatusCode,
+    Struct,
+    StructStub,
+    Tracer,
+    TypeAdapter,
+    UnsetType,
+    attrs_asdict,
+    attrs_define,
+    attrs_field,
+    attrs_fields,
+    attrs_has,
+    cattrs_structure,
+    cattrs_unstructure,
+    convert,
+    module_available,
+    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
+
+
+class ForeignKeyMetadata:
+    """Metadata for a foreign key constraint."""
+
+    __slots__ = (
+        "column_name",
+        "constraint_name",
+        "referenced_column",
+        "referenced_schema",
+        "referenced_table",
+        "schema",
+        "table_name",
+    )
+
+    def __init__(
+        self,
+        table_name: str,
+        column_name: str,
+        referenced_table: str,
+        referenced_column: str,
+        constraint_name: str | None = None,
+        schema: str | None = None,
+        referenced_schema: str | None = None,
+    ) -> None:
+        self.table_name = table_name
+        self.column_name = column_name
+        self.referenced_table = referenced_table
+        self.referenced_column = referenced_column
+        self.constraint_name = constraint_name
+        self.schema = schema
+        self.referenced_schema = referenced_schema
+
+    def __repr__(self) -> str:
+        return (
+            f"ForeignKeyMetadata(table_name={self.table_name!r}, column_name={self.column_name!r}, "
+            f"referenced_table={self.referenced_table!r}, referenced_column={self.referenced_column!r}, "
+            f"constraint_name={self.constraint_name!r}, schema={self.schema!r}, "
+            f"referenced_schema={self.referenced_schema!r})"
+        )
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, ForeignKeyMetadata):
+            return NotImplemented
+        return (
+            self.table_name == other.table_name
+            and self.column_name == other.column_name
+            and self.referenced_table == other.referenced_table
+            and self.referenced_column == other.referenced_column
+            and self.constraint_name == other.constraint_name
+            and self.schema == other.schema
+            and self.referenced_schema == other.referenced_schema
+        )
+
+    def __hash__(self) -> int:
+        return hash((
+            self.table_name,
+            self.column_name,
+            self.referenced_table,
+            self.referenced_column,
+            self.constraint_name,
+            self.schema,
+            self.referenced_schema,
+        ))
+
+
+class ColumnMetadata(TypedDict, total=False):
+    """Metadata for a database column."""
+
+    schema_name: str
+    table_name: str
+    column_name: str
+    data_type: str
+    is_nullable: str | bool | None
+    column_default: str | None
+    ordinal_position: int
+    max_length: int
+    numeric_precision: int
+    numeric_scale: int
+    is_primary: bool | int
+    is_unique: bool | int
+    extra: str
+
+
+class TableMetadata(TypedDict, total=False):
+    """Metadata for a database table."""
+
+    schema_name: str
+    table_name: str
+    table_type: str
+    table_catalog: str
+    table_schema: str
+    dependency_level: int
+    level: int
+
+
+class IndexMetadata(TypedDict, total=False):
+    """Metadata for a database index."""
+
+    schema_name: str
+    table_name: str
+    index_name: str
+    columns: list[str] | str | None
+    is_unique: bool | int
+    is_primary: bool | int
+
+
+class VersionInfo:
+    """Parsed database version info."""
+
+    def __init__(self, major: int, minor: int = 0, patch: int = 0) -> None:
+        """Initialize version info.
+
+        Args:
+            major: Major version number
+            minor: Minor version number
+            patch: Patch version number
+        """
+        self.major = major
+        self.minor = minor
+        self.patch = patch
+
+    @property
+    def version_tuple(self) -> "tuple[int, int, int]":
+        """Get version as tuple for comparison."""
+        return (self.major, self.minor, self.patch)
+
+    def __str__(self) -> str:
+        """String representation of version info."""
+        return f"{self.major}.{self.minor}.{self.patch}"
+
+    def __repr__(self) -> str:
+        """Detailed string representation."""
+        return f"VersionInfo({self.major}, {self.minor}, {self.patch})"
+
+    def __eq__(self, other: object) -> bool:
+        """Check version equality."""
+        if not isinstance(other, VersionInfo):
+            return NotImplemented
+        return self.version_tuple == other.version_tuple
+
+    def __lt__(self, other: "VersionInfo") -> bool:
+        """Check if this version is less than another."""
+        return self.version_tuple < other.version_tuple
+
+    def __le__(self, other: "VersionInfo") -> bool:
+        """Check if this version is less than or equal to another."""
+        return self.version_tuple <= other.version_tuple
+
+    def __gt__(self, other: "VersionInfo") -> bool:
+        """Check if this version is greater than another."""
+        return self.version_tuple > other.version_tuple
+
+    def __ge__(self, other: "VersionInfo") -> bool:
+        """Check if this version is greater than or equal to another."""
+        return self.version_tuple >= other.version_tuple
+
+    def __hash__(self) -> int:
+        """Make VersionInfo hashable based on version tuple."""
+        return hash(self.version_tuple)
+
+
+VersionCacheResult: TypeAlias = "tuple[bool, VersionInfo | None]"
+"""Return type for version cache lookup methods.
+
+The tuple contains:
+
+- First element (``bool``): Whether a cache lookup was attempted
+
+  - ``True``: A lookup was attempted; check second element for result
+  - ``False``: No lookup was attempted yet; second element is always ``None``
+
+- Second element (``VersionInfo | None``): The cached version info
+
+  - ``VersionInfo``: Successfully detected and cached version
+  - ``None``: Version not yet fetched, or detection failed
+"""
+
+
+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`
+"""
+SchemaT = TypeVar("SchemaT", default=dict[str, Any])
+"""Type variable for schema types (models, TypedDict, dataclasses, etc.).
+
+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.
+"""
+
+
+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 = "dict[str, object] | list[object] | tuple[object, ...] | object | None"
+"""Type alias for statement parameters.
+
+Represents:
+- :type:`dict[str, object]`
+- :type:`list[object]`
+- :type:`tuple[object, ...]`
+- :type:`object`
+- :type:`None`
+"""
+ArrowReturnFormat: TypeAlias = Literal["table", "reader", "batch", "batches"]
+"""Type alias for Apache Arrow return format options.
+
+Represents:
+- :literal:`"table"` - Return PyArrow Table
+- :literal:`"reader"` - Return PyArrow RecordBatchReader
+- :literal:`"batch"` - Return single PyArrow RecordBatch
+- :literal:`"batches"` - Return list of PyArrow RecordBatches
+"""
+
+
+@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)
+
+
+__all__ = (
+    "ALLOYDB_CONNECTOR_INSTALLED",
+    "ATTRS_INSTALLED",
+    "CATTRS_INSTALLED",
+    "CLOUD_SQL_CONNECTOR_INSTALLED",
+    "FSSPEC_INSTALLED",
+    "LITESTAR_INSTALLED",
+    "MSGSPEC_INSTALLED",
+    "NANOID_INSTALLED",
+    "NUMPY_INSTALLED",
+    "OBSTORE_INSTALLED",
+    "OPENTELEMETRY_INSTALLED",
+    "ORJSON_INSTALLED",
+    "PANDAS_INSTALLED",
+    "PGVECTOR_INSTALLED",
+    "POLARS_INSTALLED",
+    "PROMETHEUS_INSTALLED",
+    "PYARROW_INSTALLED",
+    "PYDANTIC_INSTALLED",
+    "PYDANTIC_USE_FAILFAST",
+    "UNSET",
+    "UUID_UTILS_INSTALLED",
+    "ArrowRecordBatch",
+    "ArrowRecordBatchReader",
+    "ArrowRecordBatchReaderProtocol",
+    "ArrowReturnFormat",
+    "ArrowSchema",
+    "ArrowSchemaProtocol",
+    "ArrowTable",
+    "AttrsInstance",
+    "BaseModel",
+    "ColumnMetadata",
+    "ConnectionT",
+    "Counter",
+    "DTOData",
+    "DataclassProtocol",
+    "DictLike",
+    "Empty",
+    "EmptyEnum",
+    "EmptyType",
+    "FailFast",
+    "ForeignKeyMetadata",
+    "Gauge",
+    "Histogram",
+    "IndexMetadata",
+    "NumpyArray",
+    "PandasDataFrame",
+    "PolarsDataFrame",
+    "PoolT",
+    "SchemaT",
+    "Span",
+    "StatementParameters",
+    "Status",
+    "StatusCode",
+    "Struct",
+    "SupportedSchemaModel",
+    "TableMetadata",
+    "Tracer",
+    "TypeAdapter",
+    "UnsetType",
+    "VersionCacheResult",
+    "VersionInfo",
+    "attrs_asdict",
+    "attrs_define",
+    "attrs_field",
+    "attrs_fields",
+    "attrs_has",
+    "cattrs_structure",
+    "cattrs_unstructure",
+    "convert",
+    "get_type_adapter",
+    "module_available",
+    "trace",
+)

sqlspec/utils/__init__.py

@@ -0,0 +1,7 @@
+"""Utility functions and classes for SQLSpec.
+
+This package provides various utility modules for deprecation handling,
+fixture loading, logging, module loading (including dependency checking),
+portal pattern for async bridging, singleton patterns, sync/async conversion,
+text processing, and type guards.
+"""