sqlspec 0.26.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

sqlspec/protocols.py

@@ -4,7 +4,7 @@ This module provides protocols that can be used for static type checking
 and runtime isinstance() checks.
 """
 
-from typing import TYPE_CHECKING, Any, Optional, Protocol, Union, runtime_checkable
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
 from typing_extensions import Self
 
@@ -39,6 +39,7 @@ __all__ = (
     "ParameterValueProtocol",
     "SQLBuilderProtocol",
     "SelectBuilderProtocol",
+    "SupportsArrowResults",
     "WithMethodProtocol",
 )
 
@@ -186,7 +187,7 @@ class ObjectStoreItemProtocol(Protocol):
     """Protocol for object store items with path/key attributes."""
 
     path: str
-    key: "Optional[str]"
+    key: "str | None"
 
 
 @runtime_checkable
@@ -199,35 +200,35 @@ class ObjectStoreProtocol(Protocol):
     def __init__(self, uri: str, **kwargs: Any) -> None:
         return
 
-    def read_bytes(self, path: "Union[str, Path]", **kwargs: Any) -> bytes:
+    def read_bytes(self, path: "str | Path", **kwargs: Any) -> bytes:
         """Read bytes from an object."""
         return b""
 
-    def write_bytes(self, path: "Union[str, Path]", data: bytes, **kwargs: Any) -> None:
+    def write_bytes(self, path: "str | Path", data: bytes, **kwargs: Any) -> None:
         """Write bytes to an object."""
         return
 
-    def read_text(self, path: "Union[str, Path]", encoding: str = "utf-8", **kwargs: Any) -> str:
+    def read_text(self, path: "str | Path", encoding: str = "utf-8", **kwargs: Any) -> str:
         """Read text from an object."""
         return ""
 
-    def write_text(self, path: "Union[str, Path]", data: str, encoding: str = "utf-8", **kwargs: Any) -> None:
+    def write_text(self, path: "str | Path", data: str, encoding: str = "utf-8", **kwargs: Any) -> None:
         """Write text to an object."""
         return
 
-    def exists(self, path: "Union[str, Path]", **kwargs: Any) -> bool:
+    def exists(self, path: "str | Path", **kwargs: Any) -> bool:
         """Check if an object exists."""
         return False
 
-    def delete(self, path: "Union[str, Path]", **kwargs: Any) -> None:
+    def delete(self, path: "str | Path", **kwargs: Any) -> None:
         """Delete an object."""
         return
 
-    def copy(self, source: "Union[str, Path]", destination: "Union[str, Path]", **kwargs: Any) -> None:
+    def copy(self, source: "str | Path", destination: "str | Path", **kwargs: Any) -> None:
         """Copy an object."""
         return
 
-    def move(self, source: "Union[str, Path]", destination: "Union[str, Path]", **kwargs: Any) -> None:
+    def move(self, source: "str | Path", destination: "str | Path", **kwargs: Any) -> None:
         """Move an object."""
         return
 
@@ -239,24 +240,24 @@ class ObjectStoreProtocol(Protocol):
         """Find objects matching a glob pattern."""
         return []
 
-    def is_object(self, path: "Union[str, Path]") -> bool:
+    def is_object(self, path: "str | Path") -> bool:
         """Check if path points to an object."""
         return False
 
-    def is_path(self, path: "Union[str, Path]") -> bool:
+    def is_path(self, path: "str | Path") -> bool:
         """Check if path points to a prefix (directory-like)."""
         return False
 
-    def get_metadata(self, path: "Union[str, Path]", **kwargs: Any) -> dict[str, Any]:
+    def get_metadata(self, path: "str | Path", **kwargs: Any) -> dict[str, Any]:
         """Get object metadata."""
         return {}
 
-    def read_arrow(self, path: "Union[str, Path]", **kwargs: Any) -> "ArrowTable":
+    def read_arrow(self, path: "str | Path", **kwargs: Any) -> "ArrowTable":
         """Read an Arrow table from storage."""
         msg = "Arrow reading not implemented"
         raise NotImplementedError(msg)
 
-    def write_arrow(self, path: "Union[str, Path]", table: "ArrowTable", **kwargs: Any) -> None:
+    def write_arrow(self, path: "str | Path", table: "ArrowTable", **kwargs: Any) -> None:
         """Write an Arrow table to storage."""
         msg = "Arrow writing not implemented"
         raise NotImplementedError(msg)
@@ -266,34 +267,32 @@ class ObjectStoreProtocol(Protocol):
         msg = "Arrow streaming not implemented"
         raise NotImplementedError(msg)
 
-    async def read_bytes_async(self, path: "Union[str, Path]", **kwargs: Any) -> bytes:
+    async def read_bytes_async(self, path: "str | Path", **kwargs: Any) -> bytes:
         """Async read bytes from an object."""
         msg = "Async operations not implemented"
         raise NotImplementedError(msg)
 
-    async def write_bytes_async(self, path: "Union[str, Path]", data: bytes, **kwargs: Any) -> None:
+    async def write_bytes_async(self, path: "str | Path", data: bytes, **kwargs: Any) -> None:
         """Async write bytes to an object."""
         msg = "Async operations not implemented"
         raise NotImplementedError(msg)
 
-    async def read_text_async(self, path: "Union[str, Path]", encoding: str = "utf-8", **kwargs: Any) -> str:
+    async def read_text_async(self, path: "str | Path", encoding: str = "utf-8", **kwargs: Any) -> str:
         """Async read text from an object."""
         msg = "Async operations not implemented"
         raise NotImplementedError(msg)
 
-    async def write_text_async(
-        self, path: "Union[str, Path]", data: str, encoding: str = "utf-8", **kwargs: Any
-    ) -> None:
+    async def write_text_async(self, path: "str | Path", data: str, encoding: str = "utf-8", **kwargs: Any) -> None:
         """Async write text to an object."""
         msg = "Async operations not implemented"
         raise NotImplementedError(msg)
 
-    async def exists_async(self, path: "Union[str, Path]", **kwargs: Any) -> bool:
+    async def exists_async(self, path: "str | Path", **kwargs: Any) -> bool:
         """Async check if an object exists."""
         msg = "Async operations not implemented"
         raise NotImplementedError(msg)
 
-    async def delete_async(self, path: "Union[str, Path]", **kwargs: Any) -> None:
+    async def delete_async(self, path: "str | Path", **kwargs: Any) -> None:
         """Async delete an object."""
         msg = "Async operations not implemented"
         raise NotImplementedError(msg)
@@ -303,27 +302,27 @@ class ObjectStoreProtocol(Protocol):
         msg = "Async operations not implemented"
         raise NotImplementedError(msg)
 
-    async def copy_async(self, source: "Union[str, Path]", destination: "Union[str, Path]", **kwargs: Any) -> None:
+    async def copy_async(self, source: "str | Path", destination: "str | Path", **kwargs: Any) -> None:
         """Async copy an object."""
         msg = "Async operations not implemented"
         raise NotImplementedError(msg)
 
-    async def move_async(self, source: "Union[str, Path]", destination: "Union[str, Path]", **kwargs: Any) -> None:
+    async def move_async(self, source: "str | Path", destination: "str | Path", **kwargs: Any) -> None:
         """Async move an object."""
         msg = "Async operations not implemented"
         raise NotImplementedError(msg)
 
-    async def get_metadata_async(self, path: "Union[str, Path]", **kwargs: Any) -> dict[str, Any]:
+    async def get_metadata_async(self, path: "str | Path", **kwargs: Any) -> dict[str, Any]:
         """Async get object metadata."""
         msg = "Async operations not implemented"
         raise NotImplementedError(msg)
 
-    async def read_arrow_async(self, path: "Union[str, Path]", **kwargs: Any) -> "ArrowTable":
+    async def read_arrow_async(self, path: "str | Path", **kwargs: Any) -> "ArrowTable":
         """Async read an Arrow table from storage."""
         msg = "Async arrow reading not implemented"
         raise NotImplementedError(msg)
 
-    async def write_arrow_async(self, path: "Union[str, Path]", table: "ArrowTable", **kwargs: Any) -> None:
+    async def write_arrow_async(self, path: "str | Path", table: "ArrowTable", **kwargs: Any) -> None:
         """Async write an Arrow table to storage."""
         msg = "Async arrow writing not implemented"
         raise NotImplementedError(msg)
@@ -339,7 +338,7 @@ class HasSQLGlotExpressionProtocol(Protocol):
     """Protocol for objects with a sqlglot_expression property."""
 
     @property
-    def sqlglot_expression(self) -> "Optional[exp.Expression]":
+    def sqlglot_expression(self) -> "exp.Expression | None":
         """Return the SQLGlot expression for this object."""
         ...
 
@@ -348,16 +347,24 @@ class HasSQLGlotExpressionProtocol(Protocol):
 class HasParameterBuilderProtocol(Protocol):
     """Protocol for objects that can add parameters."""
 
-    def add_parameter(self, value: Any, name: "Optional[str]" = None) -> tuple[Any, str]:
+    def add_parameter(self, value: Any, name: "str | None" = None) -> tuple[Any, str]:
         """Add a parameter to the builder."""
         ...
 
+    def get_expression(self) -> "exp.Expression | None":
+        """Return the underlying SQLGlot expression."""
+        ...
+
+    def set_expression(self, expression: "exp.Expression") -> None:
+        """Replace the underlying SQLGlot expression."""
+        ...
+
 
 @runtime_checkable
 class HasExpressionProtocol(Protocol):
     """Protocol for objects with an _expression attribute."""
 
-    _expression: "Optional[exp.Expression]"
+    _expression: "exp.Expression | None"
 
 
 @runtime_checkable
@@ -373,21 +380,21 @@ class HasToStatementProtocol(Protocol):
 class SQLBuilderProtocol(Protocol):
     """Protocol for SQL query builders."""
 
-    _expression: "Optional[exp.Expression]"
+    _expression: "exp.Expression | None"
     _parameters: dict[str, Any]
     _parameter_counter: int
     _columns: Any  # Optional attribute for some builders
     _table: Any  # Optional attribute for some builders
     _with_ctes: Any  # Optional attribute for some builders
     dialect: Any
-    dialect_name: "Optional[str]"
+    dialect_name: "str | None"
 
     @property
     def parameters(self) -> dict[str, Any]:
         """Public access to query parameters."""
         ...
 
-    def add_parameter(self, value: Any, name: "Optional[str]" = None) -> tuple[Any, str]:
+    def add_parameter(self, value: Any, name: "str | None" = None) -> tuple[Any, str]:
         """Add a parameter to the builder."""
         ...
 
@@ -399,14 +406,77 @@ class SQLBuilderProtocol(Protocol):
         """Replace literal values in an expression with bound parameters."""
         ...
 
-    def build(self) -> "Union[exp.Expression, Any]":
+    def build(self) -> "exp.Expression | Any":
         """Build and return the final expression."""
         ...
 
+    def _merge_sql_object_parameters(self, sql_obj: Any) -> None:
+        """Merge parameters from SQL objects into the builder."""
+        ...
+
+    def _build_final_expression(self, *, copy: bool = False) -> "exp.Expression":
+        """Return the expression with attached CTEs."""
+        ...
+
+    def _spawn_like_self(self) -> "Self":
+        """Create a new builder with matching configuration."""
+        ...
+
+    def get_expression(self) -> "exp.Expression | None":
+        """Return the underlying SQLGlot expression."""
+        ...
+
+    def set_expression(self, expression: "exp.Expression") -> None:
+        """Replace the underlying SQLGlot expression."""
+        ...
+
+    def generate_unique_parameter_name(self, base_name: str) -> str:
+        """Generate a unique parameter name exposed via public API."""
+        ...
+
 
 class SelectBuilderProtocol(SQLBuilderProtocol, Protocol):
     """Protocol for SELECT query builders."""
 
-    def select(self, *columns: "Union[str, exp.Expression]") -> Self:
+    def select(self, *columns: "str | exp.Expression") -> Self:
         """Add SELECT columns to the query."""
         ...
+
+
+@runtime_checkable
+class SupportsArrowResults(Protocol):
+    """Protocol for adapters that support Arrow result format.
+
+    Adapters implementing this protocol can return query results in Apache Arrow
+    format via the select_to_arrow() method, enabling zero-copy data transfer and
+    efficient integration with data science tools.
+    """
+
+    def select_to_arrow(
+        self,
+        statement: Any,
+        /,
+        *parameters: Any,
+        statement_config: Any | None = None,
+        return_format: str = "table",
+        native_only: bool = False,
+        batch_size: int | None = None,
+        arrow_schema: Any | None = None,
+        **kwargs: Any,
+    ) -> "ArrowTable | ArrowRecordBatch":
+        """Execute query and return results as Apache Arrow Table or RecordBatch.
+
+        Args:
+            statement: SQL statement to execute.
+            *parameters: Query parameters and filters.
+            statement_config: Optional statement configuration override.
+            return_format: Output format - "table", "reader", or "batches".
+            native_only: If True, raise error when native Arrow path unavailable.
+            batch_size: Chunk size for streaming modes.
+            arrow_schema: Optional target Arrow schema for type casting.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            ArrowResult containing Arrow data.
+        """
+        ...

sqlspec/storage/_utils.py

@@ -0,0 +1,85 @@
+"""Shared utilities for storage backends."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+__all__ = ("resolve_storage_path",)
+
+
+def resolve_storage_path(
+    path: "str | Path", base_path: str = "", protocol: str = "file", strip_file_scheme: bool = True
+) -> str:
+    """Resolve path relative to base_path with protocol-specific handling.
+
+    Args:
+        path: Path to resolve (may include file:// scheme).
+        base_path: Base path to prepend if path is relative.
+        protocol: Storage protocol (file, s3, gs, etc.).
+        strip_file_scheme: Whether to strip file:// prefix.
+
+    Returns:
+        Resolved path string suitable for the storage backend.
+
+    Examples:
+        >>> resolve_storage_path("/data/file.txt", protocol="file")
+        'data/file.txt'
+
+        >>> resolve_storage_path(
+        ...     "file.txt", base_path="/base", protocol="file"
+        ... )
+        'base/file.txt'
+
+        >>> resolve_storage_path(
+        ...     "file:///data/file.txt", strip_file_scheme=True
+        ... )
+        'data/file.txt'
+
+        >>> resolve_storage_path(
+        ...     "/data/subdir/file.txt",
+        ...     base_path="/data",
+        ...     protocol="file",
+        ... )
+        'subdir/file.txt'
+    """
+    from pathlib import Path as PathlibPath
+
+    path_str = str(path)
+
+    if strip_file_scheme and path_str.startswith("file://"):
+        path_str = path_str.removeprefix("file://")
+
+    # For local file protocol
+    if protocol == "file":
+        path_obj = PathlibPath(path_str)
+
+        # Absolute path handling
+        if path_obj.is_absolute():
+            if base_path:
+                base_obj = PathlibPath(base_path)
+                # Try to make path relative to base_path
+                try:
+                    relative = path_obj.relative_to(base_obj)
+                    # Return joined path for FSSpec-style backends
+                    return f"{base_path.rstrip('/')}/{relative}"
+                except ValueError:
+                    # Path is outside base_path
+                    return path_str.lstrip("/")
+            # No base_path - strip leading /
+            return path_str.lstrip("/")
+
+        # Relative path with base_path - join them
+        if base_path:
+            return f"{base_path.rstrip('/')}/{path_str}"
+
+        # Relative path without base_path
+        return path_str
+
+    # For cloud storage protocols (s3, gs, etc.), join with base_path
+    if not base_path:
+        return path_str
+
+    clean_base = base_path.rstrip("/")
+    clean_path = path_str.lstrip("/")
+    return f"{clean_base}/{clean_path}"