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

sqlspec/statement/sql_compiler.py

@@ -0,0 +1,140 @@
+"""SQL compilation logic separated from the main SQL class."""
+
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
+
+import sqlglot.expressions as exp
+
+from sqlspec.exceptions import SQLCompilationError
+from sqlspec.statement.parameters import ParameterConverter, ParameterStyle
+from sqlspec.statement.pipelines import SQLProcessingContext, StatementPipeline
+from sqlspec.statement.sql import SQLConfig
+from sqlspec.utils.cached_property import CachedProperty
+
+if TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
+
+    from sqlspec.protocols import ProcessorProtocol
+    from sqlspec.statement.parameter_manager import ParameterManager
+
+
+__all__ = ("SQLCompiler",)
+
+
+class SQLCompiler:
+    """Handles SQL compilation and pipeline processing."""
+
+    def __init__(
+        self,
+        expression: exp.Expression,
+        dialect: "Optional[DialectType]" = None,
+        parameter_manager: "Optional[ParameterManager]" = None,
+        is_script: bool = False,
+        original_sql: Optional[str] = None,
+        config: Optional[SQLConfig] = None,
+    ) -> None:
+        self.expression = expression
+        self.dialect = dialect
+        self.parameter_manager = parameter_manager
+        self.is_script = is_script
+        self._original_sql = original_sql
+        self.config = config or SQLConfig(dialect=dialect)
+
+    @CachedProperty
+    def _pipeline(self) -> StatementPipeline:
+        """Get the statement pipeline."""
+        validators: list[ProcessorProtocol] = []
+
+        if self.config.enable_validation and self.config.allowed_parameter_styles is not None:
+            from sqlspec.statement.pipelines.validators._parameter_style import ParameterStyleValidator
+
+            # In strict mode, fail on violations
+            validators.append(ParameterStyleValidator(fail_on_violation=self.config.strict_mode))
+
+        return StatementPipeline(validators=validators)
+
+    @CachedProperty
+    def _context(self) -> SQLProcessingContext:
+        """Get the processing context."""
+        if isinstance(self.expression, exp.Anonymous) and self.expression.this:
+            sql_string = str(self.expression.this)
+        else:
+            sql_string = self.expression.sql(dialect=self.dialect)
+
+        context = SQLProcessingContext(initial_sql_string=sql_string, dialect=self.dialect, config=self.config)
+        context.initial_expression = self.expression
+        context.current_expression = self.expression
+
+        from sqlspec.statement.parameters import ParameterValidator
+
+        validator = ParameterValidator()
+        context.parameter_info = validator.extract_parameters(sql_string)
+
+        if self.parameter_manager:
+            if self.parameter_manager.positional_parameters:
+                context.merged_parameters = self.parameter_manager.positional_parameters
+                context.initial_parameters = self.parameter_manager.positional_parameters
+            elif self.parameter_manager.named_parameters:
+                context.merged_parameters = self.parameter_manager.named_parameters
+                context.initial_kwargs = self.parameter_manager.named_parameters
+            context.initial_parameters = self.parameter_manager.positional_parameters
+            context.initial_kwargs = self.parameter_manager.named_parameters
+        return context
+
+    @CachedProperty
+    def _processed_expr(self) -> exp.Expression:
+        """Execute the processing pipeline and cache the result."""
+        try:
+            result = self._pipeline.execute_pipeline(self._context)
+        except Exception as e:
+            msg = f"Failed to compile SQL: {self._context.initial_sql_string}"
+            raise SQLCompilationError(msg) from e
+        else:
+            return cast("exp.Expression", result.expression)
+
+    @CachedProperty
+    def _compiled_sql(self) -> str:
+        """Get the compiled SQL string."""
+        if self.is_script:
+            return str(self._original_sql or self.expression.sql(dialect=self.dialect))
+        # Always go through the pipeline to ensure validation runs
+        processed = self._processed_expr
+        if isinstance(processed, exp.Anonymous) and processed.this:
+            return str(processed.this)
+        return str(processed.sql(dialect=self.dialect, comments=False))
+
+    def compile(self, placeholder_style: Optional[str] = None) -> tuple[str, Any]:
+        """Compile SQL and parameters."""
+        if self.is_script:
+            return self._compiled_sql, None
+
+        sql = self.to_sql(placeholder_style)
+        params = self._get_compiled_parameters(placeholder_style)
+        return sql, params
+
+    def to_sql(self, placeholder_style: Optional[str] = None) -> str:
+        """Get the SQL string with a specific placeholder style."""
+        if placeholder_style is None or self.is_script:
+            return cast("str", self._compiled_sql)
+
+        converter = ParameterConverter()
+        sql = self._compiled_sql
+
+        target_style = ParameterStyle(placeholder_style)
+        return converter.convert_placeholders(sql, target_style, self._context.parameter_info)
+
+    def get_parameters(self, style: Union[ParameterStyle, str, None] = None) -> Any:
+        """Get the parameters in a specific style."""
+        if self.is_script:
+            return None
+        return cast("Any", self._get_compiled_parameters(str(style) if style else None))
+
+    def _get_compiled_parameters(self, placeholder_style: Optional[str]) -> Any:
+        """Get compiled parameters in target style."""
+        if not self.parameter_manager:
+            return None
+
+        # This ensures the pipeline has run and context is populated
+        _ = self._processed_expr
+
+        style_enum = ParameterStyle(placeholder_style) if placeholder_style else ParameterStyle.NAMED_COLON
+        return self.parameter_manager.get_compiled_parameters(self._context.parameter_info, style_enum)

sqlspec/storage/__init__.py

@@ -5,11 +5,19 @@ This module provides a flexible storage system with:
 - Lazy loading and configuration-based registration
 - URI scheme-based automatic backend resolution
 - Key-based named storage configurations
+- Capability-based backend selection
 """
 
-from sqlspec.storage.protocol import ObjectStoreProtocol
+from sqlspec.protocols import ObjectStoreProtocol
+from sqlspec.storage.capabilities import HasStorageCapabilities, StorageCapabilities
 from sqlspec.storage.registry import StorageRegistry
 
 storage_registry = StorageRegistry()
 
-__all__ = ("ObjectStoreProtocol", "StorageRegistry", "storage_registry")
+__all__ = (
+    "HasStorageCapabilities",
+    "ObjectStoreProtocol",
+    "StorageCapabilities",
+    "StorageRegistry",
+    "storage_registry",
+)

sqlspec/storage/backends/fsspec.py

@@ -2,10 +2,11 @@
 import logging
 from io import BytesIO
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Union
 
 from sqlspec.exceptions import MissingDependencyError
 from sqlspec.storage.backends.base import ObjectStoreBase
+from sqlspec.storage.capabilities import StorageCapabilities
 from sqlspec.typing import FSSPEC_INSTALLED, PYARROW_INSTALLED
 from sqlspec.utils.sync_tools import async_
 
@@ -47,6 +48,16 @@ class FSSpecBackend(ObjectStoreBase):
     and offering fallback capabilities.
     """
 
+    # FSSpec supports most operations but varies by underlying filesystem
+    _default_capabilities: ClassVar[StorageCapabilities] = StorageCapabilities(
+        supports_arrow=PYARROW_INSTALLED,
+        supports_streaming=PYARROW_INSTALLED,
+        supports_async=True,
+        supports_compression=True,
+        is_remote=True,
+        is_cloud_native=False,
+    )
+
     def __init__(self, fs: "Union[str, AbstractFileSystem]", base_path: str = "") -> None:
         if not FSSPEC_INSTALLED:
             raise MissingDependencyError(package="fsspec", install_package="fsspec")
@@ -63,6 +74,10 @@ class FSSpecBackend(ObjectStoreBase):
             self.fs = fs
             self.protocol = getattr(fs, "protocol", "unknown")
             self._fs_uri = f"{self.protocol}://"
+
+        # Set instance-level capabilities based on detected protocol
+        self._instance_capabilities = self._detect_capabilities()
+
         super().__init__()
 
     @classmethod
@@ -71,7 +86,6 @@ class FSSpecBackend(ObjectStoreBase):
         fs_config = config.get("fs_config", {})
         base_path = config.get("base_path", "")
 
-        # Create filesystem instance from protocol
         import fsspec
 
         fs_instance = fsspec.filesystem(protocol, **fs_config)
@@ -82,12 +96,47 @@ class FSSpecBackend(ObjectStoreBase):
         """Resolve path relative to base_path."""
         path_str = str(path)
         if self.base_path:
-            # Ensure no double slashes
             clean_base = self.base_path.rstrip("/")
             clean_path = path_str.lstrip("/")
             return f"{clean_base}/{clean_path}"
         return path_str
 
+    def _detect_capabilities(self) -> StorageCapabilities:
+        """Detect capabilities based on underlying filesystem protocol."""
+        protocol = self.protocol.lower()
+
+        if protocol in {"s3", "s3a", "s3n"}:
+            return StorageCapabilities.s3_compatible()
+        if protocol in {"gcs", "gs"}:
+            return StorageCapabilities.gcs()
+        if protocol in {"abfs", "az", "azure"}:
+            return StorageCapabilities.azure_blob()
+        if protocol in {"file", "local"}:
+            return StorageCapabilities.local_filesystem()
+        return StorageCapabilities(
+            supports_arrow=PYARROW_INSTALLED,
+            supports_streaming=PYARROW_INSTALLED,
+            supports_async=True,
+            supports_compression=True,
+            is_remote=True,
+            is_cloud_native=False,
+        )
+
+    @property
+    def capabilities(self) -> StorageCapabilities:
+        """Return instance-specific capabilities based on detected protocol."""
+        return getattr(self, "_instance_capabilities", self.__class__._default_capabilities)
+
+    @classmethod
+    def has_capability(cls, capability: str) -> bool:
+        """Check if backend has a specific capability."""
+        return getattr(cls._default_capabilities, capability, False)
+
+    @classmethod
+    def get_capabilities(cls) -> StorageCapabilities:
+        """Get all capabilities for this backend."""
+        return cls._default_capabilities
+
     @property
     def backend_type(self) -> str:
         return "fsspec"
@@ -174,7 +223,6 @@ class FSSpecBackend(ObjectStoreBase):
         else:
             pattern = f"{resolved_prefix}/*" if resolved_prefix else "*"
 
-        # Get all files (not directories)
         paths = [str(path) for path in self.fs.glob(pattern, **kwargs) if not self.fs.isdir(path)]
         return sorted(paths)
 
@@ -200,7 +248,6 @@ class FSSpecBackend(ObjectStoreBase):
         """Get object metadata."""
         info = self.fs.info(self._resolve_path(path), **kwargs)
 
-        # Convert fsspec info to dict
         if isinstance(info, dict):
             return info
 
@@ -210,7 +257,6 @@ class FSSpecBackend(ObjectStoreBase):
         except AttributeError:
             pass
 
-        # Fallback to basic metadata with safe attribute access
         resolved_path = self._resolve_path(path)
         return {
             "path": resolved_path,
@@ -241,7 +287,7 @@ class FSSpecBackend(ObjectStoreBase):
         return await async_(self.read_bytes)(path, **kwargs)
 
     async def write_bytes_async(self, path: Union[str, Path], data: bytes, **kwargs: Any) -> None:
-        """Async write bytes. Wras the sync implementation."""
+        """Async write bytes. Wraps the sync implementation."""
         return await async_(self.write_bytes)(path, data, **kwargs)
 
     async def _stream_file_batches_async(self, obj_path: Union[str, Path]) -> "AsyncIterator[ArrowRecordBatch]":
@@ -268,7 +314,6 @@ class FSSpecBackend(ObjectStoreBase):
         if not PYARROW_INSTALLED:
             raise MissingDependencyError(package="pyarrow", install_package="pyarrow")
 
-        # Get paths asynchronously
         paths = await async_(self.glob)(pattern, **kwargs)
 
         # Stream batches from each path

sqlspec/storage/backends/obstore.py

@@ -9,10 +9,11 @@ from __future__ import annotations
 
 import fnmatch
 import logging
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from sqlspec.exceptions import MissingDependencyError, StorageOperationFailedError
 from sqlspec.storage.backends.base import ObjectStoreBase
+from sqlspec.storage.capabilities import HasStorageCapabilities, StorageCapabilities
 from sqlspec.typing import OBSTORE_INSTALLED
 
 if TYPE_CHECKING:
@@ -26,7 +27,7 @@ __all__ = ("ObStoreBackend",)
 logger = logging.getLogger(__name__)
 
 
-class ObStoreBackend(ObjectStoreBase):
+class ObStoreBackend(ObjectStoreBase, HasStorageCapabilities):
     """High-performance object storage backend using obstore.
 
     This backend leverages obstore's Rust-based implementation for maximum
@@ -40,6 +41,18 @@ class ObStoreBackend(ObjectStoreBase):
     Features native Arrow support and ~9x better performance than fsspec.
     """
 
+    # ObStore has excellent native capabilities
+    capabilities: ClassVar[StorageCapabilities] = StorageCapabilities(
+        supports_arrow=True,
+        supports_streaming=True,
+        supports_async=True,
+        supports_batch_operations=True,
+        supports_multipart_upload=True,
+        supports_compression=True,
+        is_cloud_native=True,
+        has_low_latency=True,
+    )
+
     def __init__(self, store_uri: str, base_path: str = "", **store_options: Any) -> None:
         """Initialize obstore backend.
 
@@ -58,14 +71,12 @@ class ObStoreBackend(ObjectStoreBase):
             self.store_options = store_options
             self.store: Any  # Will be set based on store_uri
 
-            # Initialize obstore instance
             if store_uri.startswith("memory://"):
                 # MemoryStore doesn't use from_url - create directly
                 from obstore.store import MemoryStore
 
                 self.store = MemoryStore()
             elif store_uri.startswith("file://"):
-                # For file:// URIs, use LocalStore with root directory
                 from obstore.store import LocalStore
 
                 # LocalStore works with directory paths, so we use root
@@ -86,15 +97,12 @@ class ObStoreBackend(ObjectStoreBase):
 
     def _resolve_path(self, path: str | Path) -> str:
         """Resolve path relative to base_path."""
-        # Convert Path to string
         path_str = str(path)
         # For file:// URIs, the path passed in is already absolute
         if self.store_uri.startswith("file://") and path_str.startswith("/"):
-            # Remove leading slash for LocalStore (it's relative to its root)
             return path_str.lstrip("/")
 
         if self.base_path:
-            # Ensure no double slashes by stripping trailing slash from base_path
             clean_base = self.base_path.rstrip("/")
             clean_path = path_str.lstrip("/")
             return f"{clean_base}/{clean_path}"
@@ -113,7 +121,6 @@ class ObStoreBackend(ObjectStoreBase):
             resolved_path = self._resolve_path(path)
             result = self.store.get(resolved_path)
             bytes_data = result.bytes()
-            # Handle obstore's Bytes type - it might have a method to get raw bytes
             if hasattr(bytes_data, "__bytes__"):
                 return bytes(bytes_data)
             if hasattr(bytes_data, "tobytes"):
@@ -210,17 +217,14 @@ class ObStoreBackend(ObjectStoreBase):
         resolved_pattern = self._resolve_path(pattern)
         all_objects = self.list_objects(recursive=True, **kwargs)
 
-        # For complex patterns with **, use PurePosixPath
         if "**" in pattern:
             matching_objects = []
 
             # Special case: **/*.ext should also match *.ext in root
             if pattern.startswith("**/"):
-                # Get the suffix pattern
                 suffix_pattern = pattern[3:]  # Remove **/
 
                 for obj in all_objects:
-                    # Check if object ends with the suffix pattern
                     obj_path = PurePosixPath(obj)
                     # Try both the full pattern and just the suffix
                     if obj_path.match(resolved_pattern) or obj_path.match(suffix_pattern):
@@ -271,7 +275,6 @@ class ObStoreBackend(ObjectStoreBase):
         if resolved_path.endswith("/"):
             return True
 
-        # Check if there are any objects with this prefix
         try:
             objects = self.list_objects(prefix=str(path), recursive=False)
             return len(objects) > 0
@@ -282,7 +285,6 @@ class ObStoreBackend(ObjectStoreBase):
         """Read Arrow table using obstore."""
         try:
             resolved_path = self._resolve_path(path)
-            # Check if the store has native Arrow support
             if hasattr(self.store, "read_arrow"):
                 return self.store.read_arrow(resolved_path, **kwargs)  # type: ignore[no-any-return]  # pyright: ignore[reportAttributeAccessIssue]
             # Fall back to reading as Parquet via bytes
@@ -301,7 +303,6 @@ class ObStoreBackend(ObjectStoreBase):
         """Write Arrow table using obstore."""
         try:
             resolved_path = self._resolve_path(path)
-            # Check if the store has native Arrow support
             if hasattr(self.store, "write_arrow"):
                 self.store.write_arrow(resolved_path, table, **kwargs)  # pyright: ignore[reportAttributeAccessIssue]
             else:
@@ -321,7 +322,6 @@ class ObStoreBackend(ObjectStoreBase):
 
                 for field in schema:
                     if str(field.type).startswith("decimal64"):
-                        # Convert decimal64 to decimal128
                         import re
 
                         match = re.match(r"decimal64\((\d+),\s*(\d+)\)", str(field.type))
@@ -367,7 +367,6 @@ class ObStoreBackend(ObjectStoreBase):
         resolved_path = self._resolve_path(path)
         result = await self.store.get_async(resolved_path)
         bytes_data = result.bytes()
-        # Handle obstore's Bytes type - it might have a method to get raw bytes
         if hasattr(bytes_data, "__bytes__"):
             return bytes(bytes_data)
         if hasattr(bytes_data, "tobytes"):
@@ -441,10 +440,8 @@ class ObStoreBackend(ObjectStoreBase):
         resolved_path = self._resolve_path(path)
         metadata = await self.store.head_async(resolved_path)
 
-        # Convert obstore ObjectMeta to dict
         result = {"path": resolved_path, "exists": True}
 
-        # Extract metadata attributes if available
         for attr in ["size", "last_modified", "e_tag", "version"]:
             if hasattr(metadata, attr):
                 result[attr] = getattr(metadata, attr)
@@ -465,7 +462,6 @@ class ObStoreBackend(ObjectStoreBase):
     async def write_arrow_async(self, path: str | Path, table: ArrowTable, **kwargs: Any) -> None:
         """Async write Arrow table using native obstore async."""
         resolved_path = self._resolve_path(path)
-        # Check if the store has native async Arrow support
         if hasattr(self.store, "write_arrow_async"):
             await self.store.write_arrow_async(resolved_path, table, **kwargs)  # pyright: ignore[reportAttributeAccessIssue]
         else:

sqlspec/storage/capabilities.py

@@ -0,0 +1,101 @@
+"""Storage backend capability system.
+
+This module provides a centralized way to track and query storage backend capabilities.
+"""
+
+from dataclasses import dataclass
+from typing import ClassVar
+
+__all__ = ("HasStorageCapabilities", "StorageCapabilities")
+
+
+@dataclass
+class StorageCapabilities:
+    """Tracks capabilities of a storage backend."""
+
+    # Basic operations
+    supports_read: bool = True
+    supports_write: bool = True
+    supports_delete: bool = True
+    supports_list: bool = True
+    supports_exists: bool = True
+    supports_copy: bool = True
+    supports_move: bool = True
+    supports_metadata: bool = True
+
+    # Advanced operations
+    supports_arrow: bool = False
+    supports_streaming: bool = False
+    supports_async: bool = False
+    supports_batch_operations: bool = False
+    supports_multipart_upload: bool = False
+    supports_compression: bool = False
+
+    # Protocol-specific features
+    supports_s3_select: bool = False
+    supports_gcs_compose: bool = False
+    supports_azure_snapshots: bool = False
+
+    # Performance characteristics
+    is_remote: bool = True
+    is_cloud_native: bool = False
+    has_low_latency: bool = False
+
+    @classmethod
+    def local_filesystem(cls) -> "StorageCapabilities":
+        """Capabilities for local filesystem backend."""
+        return cls(
+            is_remote=False, has_low_latency=True, supports_arrow=True, supports_streaming=True, supports_async=True
+        )
+
+    @classmethod
+    def s3_compatible(cls) -> "StorageCapabilities":
+        """Capabilities for S3-compatible backends."""
+        return cls(
+            is_cloud_native=True,
+            supports_multipart_upload=True,
+            supports_s3_select=True,
+            supports_arrow=True,
+            supports_streaming=True,
+            supports_async=True,
+        )
+
+    @classmethod
+    def gcs(cls) -> "StorageCapabilities":
+        """Capabilities for Google Cloud Storage."""
+        return cls(
+            is_cloud_native=True,
+            supports_multipart_upload=True,
+            supports_gcs_compose=True,
+            supports_arrow=True,
+            supports_streaming=True,
+            supports_async=True,
+        )
+
+    @classmethod
+    def azure_blob(cls) -> "StorageCapabilities":
+        """Capabilities for Azure Blob Storage."""
+        return cls(
+            is_cloud_native=True,
+            supports_multipart_upload=True,
+            supports_azure_snapshots=True,
+            supports_arrow=True,
+            supports_streaming=True,
+            supports_async=True,
+        )
+
+
+class HasStorageCapabilities:
+    """Mixin for storage backends that expose their capabilities."""
+
+    capabilities: ClassVar[StorageCapabilities]
+
+    @classmethod
+    def has_capability(cls, capability: str) -> bool:
+        """Check if backend has a specific capability."""
+        return getattr(cls.capabilities, capability, False)
+
+    @classmethod
+    def get_capabilities(cls) -> StorageCapabilities:
+        """Get all capabilities for this backend."""
+        return cls.capabilities