sqlspec 0.12.2__py3-none-any.whl → 0.13.0__py3-none-any.whl

sqlspec/storage/registry.py

@@ -8,13 +8,13 @@ This module provides a flexible, lazy-loading storage registry that supports:
 - Automatic instrumentation integration
 """
 
-# TODO: TRY300 - Review try-except patterns for else block opportunities
 import logging
 from pathlib import Path
 from typing import Any, Optional, TypeVar, Union, cast
 
 from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError
-from sqlspec.storage.protocol import ObjectStoreProtocol
+from sqlspec.protocols import ObjectStoreProtocol
+from sqlspec.storage.capabilities import StorageCapabilities
 from sqlspec.typing import FSSPEC_INSTALLED, OBSTORE_INSTALLED
 
 __all__ = ("StorageRegistry", "storage_registry")
@@ -58,9 +58,7 @@ class StorageRegistry:
     """
 
     def __init__(self) -> None:
-        # Named aliases (secondary feature) - internal storage
         self._alias_configs: dict[str, tuple[type[ObjectStoreProtocol], str, dict[str, Any]]] = {}
-        # Expose configs for testing compatibility
         self._aliases: dict[str, dict[str, Any]] = {}
         self._instances: dict[Union[str, tuple[str, tuple[tuple[str, Any], ...]]], ObjectStoreProtocol] = {}
 
@@ -85,21 +83,17 @@ class StorageRegistry:
             **kwargs: Backend-specific configuration options
         """
         if backend is None:
-            # Auto-detect from URI using new intelligent selection
             backend = self._determine_backend_class(uri)
 
         config = config or {}
         config.update(kwargs)
 
-        # Store the actual config that will be passed to backend
         backend_config = dict(config)
         if base_path:
             backend_config["base_path"] = base_path
 
-        # Store backend class, URI, and config separately
         self._alias_configs[alias] = (backend, uri, backend_config)
 
-        # Store config with URI for test compatibility
         test_config = dict(backend_config)
         test_config["uri"] = uri
         self._aliases[alias] = test_config
@@ -117,72 +111,35 @@ class StorageRegistry:
         Raises:
             ImproperConfigurationError: If alias not found or invalid input
         """
-        # Handle None case - raise AttributeError for test compatibility
-        if uri_or_alias is None:
-            msg = "uri_or_alias cannot be None"
-            raise AttributeError(msg)
-
-        # Handle empty string
         if not uri_or_alias:
-            msg = "Unknown storage alias: ''"
+            msg = "URI or alias cannot be empty."
             raise ImproperConfigurationError(msg)
 
-        # Handle Path objects - convert to file:// URI
         if isinstance(uri_or_alias, Path):
             uri_or_alias = f"file://{uri_or_alias.resolve()}"
 
-        # Check cache first
-        cache_key: Union[str, tuple[str, tuple[tuple[str, Any], ...]]] = (
-            (uri_or_alias, tuple(sorted(kwargs.items()))) if kwargs else uri_or_alias
-        )
+        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]
 
-        # PRIMARY: Try URI-first routing
         if "://" in uri_or_alias:
-            backend = self._resolve_from_uri(uri_or_alias, **kwargs)
-            # Cache the instance for future use
-            self._instances[cache_key] = backend
-            return backend
-
-        # SECONDARY: Check if it's a registered alias
-        if uri_or_alias in self._alias_configs:
+            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]
-            # Merge kwargs with alias config (kwargs override)
-            merged_config = dict(config)
-            merged_config.update(kwargs)
-            # URI is passed as first positional arg
+            merged_config = {**config, **kwargs}
             instance = backend_cls(stored_uri, **merged_config)
-            self._instances[cache_key] = instance
-            return instance
+        else:
+            msg = f"Unknown storage alias or invalid URI: '{uri_or_alias}'"
+            raise ImproperConfigurationError(msg)
 
-        # Not a URI and not an alias
-        msg = f"Unknown storage alias: '{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.
-
-        Tries ObStore first for supported schemes, then falls back to FSSpec.
-
-        Args:
-            uri: URI to resolve backend for
-            **kwargs: Additional backend-specific configuration
-
-        Returns:
-            Backend instance
-
-        Raises:
-            MissingDependencyError: If no suitable backend can be created
-        """
-        # Schemes that ObStore doesn't support
-
-        # Extract scheme
+        """Resolve backend from URI, trying ObStore first, then FSSpec."""
         scheme = self._get_scheme(uri)
-
         last_exc: Optional[Exception] = None
 
-        # If scheme is FSSpec-only, skip ObStore
         if scheme not in FSSPEC_ONLY_SCHEMES and OBSTORE_INSTALLED:
             try:
                 return self._create_backend("obstore", uri, **kwargs)
@@ -197,13 +154,14 @@ class StorageRegistry:
                 logger.debug("FSSpec backend failed for %s: %s", uri, e)
                 last_exc = e
 
-        msg = f"No storage backend available for URI '{uri}'. Install 'obstore' or 'fsspec' and ensure dependencies for your filesystem are installed."
+        msg = f"No storage backend available for URI '{uri}'. Install 'obstore' or 'fsspec' and required dependencies."
         raise MissingDependencyError(msg) from last_exc
 
     def _determine_backend_class(self, uri: str) -> type[ObjectStoreProtocol]:
-        """Determine the best backend class for a URI based on availability.
+        """Determine the best backend class for a URI based on availability and capabilities.
 
-        Prefers ObStore, falls back to FSSpec.
+        Prefers ObStore for its superior performance and native capabilities,
+        falls back to FSSpec for extended protocol support.
 
         Args:
             uri: URI to determine backend for.
@@ -211,12 +169,20 @@ class StorageRegistry:
         Returns:
             Backend class (not instance)
         """
+        scheme = self._get_scheme(uri)
+
+        # Check if scheme requires FSSpec (not supported by ObStore)
+        if scheme in FSSPEC_ONLY_SCHEMES and FSSPEC_INSTALLED:
+            return self._get_backend_class("fsspec")
+
+        # Prefer ObStore for its superior performance
         if OBSTORE_INSTALLED:
             return self._get_backend_class("obstore")
+            # Could check capabilities here if needed
+
         if FSSPEC_INSTALLED:
             return self._get_backend_class("fsspec")
 
-        scheme = uri.split("://", maxsplit=1)[0].lower()
         msg = f"No backend available for URI scheme '{scheme}'. Install obstore or fsspec."
         raise MissingDependencyError(msg)
 
@@ -244,35 +210,14 @@ class StorageRegistry:
         raise ValueError(msg)
 
     def _create_backend(self, backend_type: str, uri: str, **kwargs: Any) -> ObjectStoreProtocol:
-        """Create backend instance for URI.
-
-        Args:
-            backend_type: Backend type ('obstore' or 'fsspec')
-            uri: URI to create backend for
-            **kwargs: Additional backend-specific configuration
-
-        Returns:
-            Backend instance
-        """
+        """Create backend instance for URI."""
         backend_cls = self._get_backend_class(backend_type)
-        # Both backends accept URI as first positional parameter
         return backend_cls(uri, **kwargs)
 
     def _get_scheme(self, uri: str) -> str:
-        """Extract scheme from URI.
-
-        Args:
-            uri: URI to extract scheme from
-
-        Returns:
-            Scheme (e.g., 's3', 'gs', 'file')
-        """
-        # Handle file paths without explicit file:// scheme
+        """Extract scheme from URI."""
         if not uri or "://" not in uri:
-            # Local path (absolute or relative)
             return "file"
-
-        # Extract scheme from URI
         return uri.split("://", maxsplit=1)[0].lower()
 
     # Utility methods
@@ -310,6 +255,34 @@ class StorageRegistry:
         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.
+
+        Args:
+            uri_or_alias: URI or alias to check capabilities for
+
+        Returns:
+            StorageCapabilities object describing backend capabilities
+        """
+        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)
+
+        # Get capabilities from the backend class
+        if hasattr(backend_cls, "capabilities"):
+            return backend_cls.capabilities
+
+        # Default capabilities if not defined
+
+        return StorageCapabilities()
+
 
 # Global registry instance
 storage_registry = StorageRegistry()