sqlspec 0.26.0__py3-none-any.whl → 0.27.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
 
@@ -186,7 +186,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 +199,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 +239,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 +266,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 +301,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 +337,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 +346,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 +379,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 +405,38 @@ 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."""
         ...

sqlspec/storage/_utils.py

@@ -0,0 +1,98 @@
+"""Shared utilities for storage backends."""
+
+from typing import TYPE_CHECKING
+
+from sqlspec.exceptions import MissingDependencyError
+from sqlspec.typing import PYARROW_INSTALLED
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+__all__ = ("ensure_pyarrow", "resolve_storage_path")
+
+
+def ensure_pyarrow() -> None:
+    """Ensure PyArrow is available for Arrow operations.
+
+    Raises:
+        MissingDependencyError: If pyarrow is not installed.
+    """
+    if not PYARROW_INSTALLED:
+        raise MissingDependencyError(package="pyarrow", install_package="pyarrow")
+
+
+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}"