sqlspec 0.12.1__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

@@ -1,10 +1,12 @@
 # pyright: ignore=reportUnknownVariableType
 import logging
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, Union
+from pathlib import Path
+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_
 
@@ -46,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")
@@ -62,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
@@ -70,21 +86,56 @@ 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)
 
         return cls(fs=fs_instance, base_path=base_path)
 
-    def _resolve_path(self, path: str) -> str:
+    def _resolve_path(self, path: Union[str, Path]) -> str:
         """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.lstrip("/")
+            clean_path = path_str.lstrip("/")
             return f"{clean_base}/{clean_path}"
-        return 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:
@@ -95,51 +146,51 @@ class FSSpecBackend(ObjectStoreBase):
         return self._fs_uri
 
     # Core Operations (sync)
-    def read_bytes(self, path: str, **kwargs: Any) -> bytes:
+    def read_bytes(self, path: Union[str, Path], **kwargs: Any) -> bytes:
         """Read bytes from an object."""
         resolved_path = self._resolve_path(path)
         return self.fs.cat(resolved_path, **kwargs)  # type: ignore[no-any-return]  # pyright: ignore
 
-    def write_bytes(self, path: str, data: bytes, **kwargs: Any) -> None:
+    def write_bytes(self, path: Union[str, Path], data: bytes, **kwargs: Any) -> None:
         """Write bytes to an object."""
         resolved_path = self._resolve_path(path)
         with self.fs.open(resolved_path, mode="wb", **kwargs) as f:
             f.write(data)  # pyright: ignore
 
-    def read_text(self, path: str, encoding: str = "utf-8", **kwargs: Any) -> str:
+    def read_text(self, path: Union[str, Path], encoding: str = "utf-8", **kwargs: Any) -> str:
         """Read text from an object."""
         data = self.read_bytes(path, **kwargs)
         return data.decode(encoding)
 
-    def write_text(self, path: str, data: str, encoding: str = "utf-8", **kwargs: Any) -> None:
+    def write_text(self, path: Union[str, Path], data: str, encoding: str = "utf-8", **kwargs: Any) -> None:
         """Write text to an object."""
         self.write_bytes(path, data.encode(encoding), **kwargs)
 
     # Object Operations
-    def exists(self, path: str, **kwargs: Any) -> bool:
+    def exists(self, path: Union[str, Path], **kwargs: Any) -> bool:
         """Check if an object exists."""
         resolved_path = self._resolve_path(path)
         return self.fs.exists(resolved_path, **kwargs)  # type: ignore[no-any-return]
 
-    def delete(self, path: str, **kwargs: Any) -> None:
+    def delete(self, path: Union[str, Path], **kwargs: Any) -> None:
         """Delete an object."""
         resolved_path = self._resolve_path(path)
         self.fs.rm(resolved_path, **kwargs)
 
-    def copy(self, source: str, destination: str, **kwargs: Any) -> None:
+    def copy(self, source: Union[str, Path], destination: Union[str, Path], **kwargs: Any) -> None:
         """Copy an object."""
         source_path = self._resolve_path(source)
         dest_path = self._resolve_path(destination)
         self.fs.copy(source_path, dest_path, **kwargs)
 
-    def move(self, source: str, destination: str, **kwargs: Any) -> None:
+    def move(self, source: Union[str, Path], destination: Union[str, Path], **kwargs: Any) -> None:
         """Move an object."""
         source_path = self._resolve_path(source)
         dest_path = self._resolve_path(destination)
         self.fs.mv(source_path, dest_path, **kwargs)
 
     # Arrow Operations
-    def read_arrow(self, path: str, **kwargs: Any) -> "ArrowTable":
+    def read_arrow(self, path: Union[str, Path], **kwargs: Any) -> "ArrowTable":
         """Read an Arrow table from storage."""
         if not PYARROW_INSTALLED:
             raise MissingDependencyError(package="pyarrow", install_package="pyarrow")
@@ -150,7 +201,7 @@ class FSSpecBackend(ObjectStoreBase):
         with self.fs.open(resolved_path, mode="rb", **kwargs) as f:
             return pq.read_table(f)
 
-    def write_arrow(self, path: str, table: "ArrowTable", **kwargs: Any) -> None:
+    def write_arrow(self, path: Union[str, Path], table: "ArrowTable", **kwargs: Any) -> None:
         """Write an Arrow table to storage."""
         if not PYARROW_INSTALLED:
             raise MissingDependencyError(package="pyarrow", install_package="pyarrow")
@@ -172,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)
 
@@ -194,11 +244,10 @@ class FSSpecBackend(ObjectStoreBase):
         resolved_path = self._resolve_path(path)
         return self.fs.isdir(resolved_path)  # type: ignore[no-any-return]
 
-    def get_metadata(self, path: str, **kwargs: Any) -> dict[str, Any]:
+    def get_metadata(self, path: Union[str, Path], **kwargs: Any) -> dict[str, Any]:
         """Get object metadata."""
         info = self.fs.info(self._resolve_path(path), **kwargs)
 
-        # Convert fsspec info to dict
         if isinstance(info, dict):
             return info
 
@@ -208,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,
@@ -217,7 +265,7 @@ class FSSpecBackend(ObjectStoreBase):
             "type": getattr(info, "type", "file"),
         }
 
-    def _stream_file_batches(self, obj_path: str) -> "Iterator[ArrowRecordBatch]":
+    def _stream_file_batches(self, obj_path: Union[str, Path]) -> "Iterator[ArrowRecordBatch]":
         import pyarrow.parquet as pq
 
         with self.fs.open(obj_path, mode="rb") as f:
@@ -234,15 +282,15 @@ class FSSpecBackend(ObjectStoreBase):
         for obj_path in self.glob(pattern, **kwargs):
             yield from self._stream_file_batches(obj_path)
 
-    async def read_bytes_async(self, path: str, **kwargs: Any) -> bytes:
+    async def read_bytes_async(self, path: Union[str, Path], **kwargs: Any) -> bytes:
         """Async read bytes. Wraps the sync implementation."""
         return await async_(self.read_bytes)(path, **kwargs)
 
-    async def write_bytes_async(self, path: str, data: bytes, **kwargs: Any) -> None:
-        """Async write bytes. Wras the sync implementation."""
+    async def write_bytes_async(self, path: Union[str, Path], data: bytes, **kwargs: Any) -> None:
+        """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: str) -> "AsyncIterator[ArrowRecordBatch]":
+    async def _stream_file_batches_async(self, obj_path: Union[str, Path]) -> "AsyncIterator[ArrowRecordBatch]":
         import pyarrow.parquet as pq
 
         data = await self.read_bytes_async(obj_path)
@@ -266,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
@@ -274,11 +321,11 @@ class FSSpecBackend(ObjectStoreBase):
             async for batch in self._stream_file_batches_async(path):
                 yield batch
 
-    async def read_text_async(self, path: str, encoding: str = "utf-8", **kwargs: Any) -> str:
+    async def read_text_async(self, path: Union[str, Path], encoding: str = "utf-8", **kwargs: Any) -> str:
         """Async read text. Wraps the sync implementation."""
         return await async_(self.read_text)(path, encoding, **kwargs)
 
-    async def write_text_async(self, path: str, data: str, encoding: str = "utf-8", **kwargs: Any) -> None:
+    async def write_text_async(self, path: Union[str, Path], data: str, encoding: str = "utf-8", **kwargs: Any) -> None:
         """Async write text. Wraps the sync implementation."""
         await async_(self.write_text)(path, data, encoding, **kwargs)
 
@@ -286,30 +333,30 @@ class FSSpecBackend(ObjectStoreBase):
         """Async list objects. Wraps the sync implementation."""
         return await async_(self.list_objects)(prefix, recursive, **kwargs)
 
-    async def exists_async(self, path: str, **kwargs: Any) -> bool:
+    async def exists_async(self, path: Union[str, Path], **kwargs: Any) -> bool:
         """Async exists check. Wraps the sync implementation."""
         return await async_(self.exists)(path, **kwargs)
 
-    async def delete_async(self, path: str, **kwargs: Any) -> None:
+    async def delete_async(self, path: Union[str, Path], **kwargs: Any) -> None:
         """Async delete. Wraps the sync implementation."""
         await async_(self.delete)(path, **kwargs)
 
-    async def copy_async(self, source: str, destination: str, **kwargs: Any) -> None:
+    async def copy_async(self, source: Union[str, Path], destination: Union[str, Path], **kwargs: Any) -> None:
         """Async copy. Wraps the sync implementation."""
         await async_(self.copy)(source, destination, **kwargs)
 
-    async def move_async(self, source: str, destination: str, **kwargs: Any) -> None:
+    async def move_async(self, source: Union[str, Path], destination: Union[str, Path], **kwargs: Any) -> None:
         """Async move. Wraps the sync implementation."""
         await async_(self.move)(source, destination, **kwargs)
 
-    async def get_metadata_async(self, path: str, **kwargs: Any) -> dict[str, Any]:
+    async def get_metadata_async(self, path: Union[str, Path], **kwargs: Any) -> dict[str, Any]:
         """Async get metadata. Wraps the sync implementation."""
         return await async_(self.get_metadata)(path, **kwargs)
 
-    async def read_arrow_async(self, path: str, **kwargs: Any) -> "ArrowTable":
+    async def read_arrow_async(self, path: Union[str, Path], **kwargs: Any) -> "ArrowTable":
         """Async read Arrow. Wraps the sync implementation."""
         return await async_(self.read_arrow)(path, **kwargs)
 
-    async def write_arrow_async(self, path: str, table: "ArrowTable", **kwargs: Any) -> None:
+    async def write_arrow_async(self, path: Union[str, Path], table: "ArrowTable", **kwargs: Any) -> None:
         """Async write Arrow. Wraps the sync implementation."""
         await async_(self.write_arrow)(path, table, **kwargs)