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

sqlspec/storage/backends/fsspec.py

@@ -1,10 +1,13 @@
 # pyright: reportPrivateUsage=false
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any
+
+from mypy_extensions import mypyc_attr
 
 from sqlspec.exceptions import MissingDependencyError
-from sqlspec.typing import FSSPEC_INSTALLED, PYARROW_INSTALLED
+from sqlspec.storage._utils import ensure_pyarrow, resolve_storage_path
+from sqlspec.typing import FSSPEC_INSTALLED
 from sqlspec.utils.sync_tools import async_
 
 if TYPE_CHECKING:
@@ -18,41 +21,80 @@ logger = logging.getLogger(__name__)
 
 
 class _ArrowStreamer:
+    """Async iterator for streaming Arrow batches from FSSpec backend.
+
+    Uses async_() to offload blocking operations to thread pool,
+    preventing event loop blocking during file I/O and iteration.
+
+    CRITICAL: Creates generators on main thread, offloads only next() calls.
+    """
+
+    __slots__ = ("_initialized", "backend", "batch_iterator", "kwargs", "paths_iterator", "pattern")
+
     def __init__(self, backend: "FSSpecBackend", pattern: str, **kwargs: Any) -> None:
         self.backend = backend
         self.pattern = pattern
         self.kwargs = kwargs
-        self.paths_iterator: Optional[Iterator[str]] = None
-        self.batch_iterator: Optional[Iterator[ArrowRecordBatch]] = None
+        self.paths_iterator: Iterator[str] | None = None
+        self.batch_iterator: Iterator[ArrowRecordBatch] | None = None
+        self._initialized = False
 
     def __aiter__(self) -> "_ArrowStreamer":
         return self
 
     async def _initialize(self) -> None:
-        """Initialize paths iterator."""
-        if self.paths_iterator is None:
+        """Initialize paths iterator asynchronously."""
+        if not self._initialized:
             paths = await async_(self.backend.glob)(self.pattern, **self.kwargs)
             self.paths_iterator = iter(paths)
+            self._initialized = True
 
     async def __anext__(self) -> "ArrowRecordBatch":
+        """Get next Arrow batch asynchronously.
+
+        Iterative state machine that avoids recursion and blocking calls.
+
+        Returns:
+            Arrow record batches from matching files.
+
+        Raises:
+            StopAsyncIteration: When no more batches available.
+        """
         await self._initialize()
 
-        if self.batch_iterator:
+        while True:
+            if self.batch_iterator is not None:
+
+                def _safe_next_batch() -> "ArrowRecordBatch":
+                    try:
+                        return next(self.batch_iterator)  # type: ignore[arg-type]
+                    except StopIteration as e:
+                        raise StopAsyncIteration from e
+
+                try:
+                    return await async_(_safe_next_batch)()
+                except StopAsyncIteration:
+                    self.batch_iterator = None
+                    continue
+
             try:
-                return next(self.batch_iterator)
-            except StopIteration:
-                self.batch_iterator = None
+                path = next(self.paths_iterator)  # type: ignore[arg-type]
+            except StopIteration as e:
+                raise StopAsyncIteration from e
+
+            self.batch_iterator = self.backend._stream_file_batches(path)
 
-        if self.paths_iterator:
+    async def aclose(self) -> None:
+        """Close underlying batch iterator."""
+        if self.batch_iterator is not None:
             try:
-                path = next(self.paths_iterator)
-                self.batch_iterator = await async_(self.backend._stream_file_batches)(path)
-                return await self.__anext__()
-            except StopIteration:
-                raise StopAsyncIteration
-        raise StopAsyncIteration
+                close_method = self.batch_iterator.close  # type: ignore[attr-defined]
+                await async_(close_method)()
+            except AttributeError:
+                pass
 
 
+@mypyc_attr(allow_interpreted_subclasses=True)
 class FSSpecBackend:
     """Storage backend using fsspec.
 
@@ -60,19 +102,37 @@ class FSSpecBackend:
     including HTTP, HTTPS, FTP, and cloud storage services.
     """
 
+    __slots__ = ("_fs_uri", "backend_type", "base_path", "fs", "protocol")
+
     def __init__(self, uri: str, **kwargs: Any) -> None:
-        self._ensure_fsspec()
+        if not FSSPEC_INSTALLED:
+            raise MissingDependencyError(package="fsspec", install_package="fsspec")
 
         base_path = kwargs.pop("base_path", "")
-        self.base_path = base_path.rstrip("/") if base_path else ""
 
         if "://" in uri:
             self.protocol = uri.split("://", maxsplit=1)[0]
             self._fs_uri = uri
+
+            # For S3/cloud URIs, extract bucket/path from URI as base_path
+            if self.protocol in {"s3", "gs", "az", "gcs"}:
+                from urllib.parse import urlparse
+
+                parsed = urlparse(uri)
+                # Combine netloc (bucket) and path for base_path
+                if parsed.netloc:
+                    uri_base_path = parsed.netloc
+                    if parsed.path and parsed.path != "/":
+                        uri_base_path = f"{uri_base_path}{parsed.path}"
+                    # Only use URI base_path if no explicit base_path provided
+                    if not base_path:
+                        base_path = uri_base_path
         else:
             self.protocol = uri
             self._fs_uri = f"{uri}://"
 
+        self.base_path = base_path.rstrip("/") if base_path else ""
+
         import fsspec
 
         self.fs = fsspec.filesystem(self.protocol, **kwargs)
@@ -93,48 +153,19 @@ class FSSpecBackend:
 
         return cls(uri=uri, **kwargs)
 
-    def _ensure_fsspec(self) -> None:
-        """Ensure fsspec is available for operations."""
-        if not FSSPEC_INSTALLED:
-            raise MissingDependencyError(package="fsspec", install_package="fsspec")
-
-    def _ensure_pyarrow(self) -> None:
-        """Ensure PyArrow is available for Arrow operations."""
-        if not PYARROW_INSTALLED:
-            raise MissingDependencyError(package="pyarrow", install_package="pyarrow")
-
-    def _resolve_path(self, path: Union[str, Path]) -> str:
-        """Resolve path relative to base_path."""
-        path_str = str(path)
-        if self.base_path:
-            clean_base = self.base_path.rstrip("/")
-            clean_path = path_str.lstrip("/")
-            return f"{clean_base}/{clean_path}"
-        if self.protocol == "s3" and "://" in self._fs_uri:
-            # For S3, we need to include the bucket from the URI
-            # Extract bucket and path from URI like s3://bucket/path
-            uri_parts = self._fs_uri.split("://", 1)[1]  # Remove s3://
-            if "/" in uri_parts:
-                # URI has bucket and base path
-                return f"{uri_parts.rstrip('/')}/{path_str.lstrip('/')}"
-            # URI has only bucket
-            return f"{uri_parts}/{path_str.lstrip('/')}"
-        return path_str
-
     @property
     def base_uri(self) -> str:
         return self._fs_uri
 
-    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."""
-        resolved_path = self._resolve_path(path)
+        resolved_path = resolve_storage_path(path, self.base_path, self.protocol, strip_file_scheme=False)
         return self.fs.cat(resolved_path, **kwargs)  # type: ignore[no-any-return]  # pyright: ignore
 
-    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."""
-        resolved_path = self._resolve_path(path)
+        resolved_path = resolve_storage_path(path, self.base_path, self.protocol, strip_file_scheme=False)
 
-        # Only create directories for local file systems, not for cloud storage
         if self.protocol == "file":
             parent_dir = str(Path(resolved_path).parent)
             if parent_dir and not self.fs.exists(parent_dir):
@@ -143,84 +174,85 @@ class FSSpecBackend:
         with self.fs.open(resolved_path, mode="wb", **kwargs) as f:
             f.write(data)  # pyright: ignore
 
-    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."""
         data = self.read_bytes(path, **kwargs)
         return data.decode(encoding)
 
-    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."""
         self.write_bytes(path, data.encode(encoding), **kwargs)
 
-    def exists(self, path: Union[str, Path], **kwargs: Any) -> bool:
+    def exists(self, path: str | Path, **kwargs: Any) -> bool:
         """Check if an object exists."""
-        resolved_path = self._resolve_path(path)
+        resolved_path = resolve_storage_path(path, self.base_path, self.protocol, strip_file_scheme=False)
         return self.fs.exists(resolved_path, **kwargs)  # type: ignore[no-any-return]
 
-    def delete(self, path: Union[str, Path], **kwargs: Any) -> None:
+    def delete(self, path: str | Path, **kwargs: Any) -> None:
         """Delete an object."""
-        resolved_path = self._resolve_path(path)
+        resolved_path = resolve_storage_path(path, self.base_path, self.protocol, strip_file_scheme=False)
         self.fs.rm(resolved_path, **kwargs)
 
-    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."""
-        source_path = self._resolve_path(source)
-        dest_path = self._resolve_path(destination)
+        source_path = resolve_storage_path(source, self.base_path, self.protocol, strip_file_scheme=False)
+        dest_path = resolve_storage_path(destination, self.base_path, self.protocol, strip_file_scheme=False)
         self.fs.copy(source_path, dest_path, **kwargs)
 
-    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."""
-        source_path = self._resolve_path(source)
-        dest_path = self._resolve_path(destination)
+        source_path = resolve_storage_path(source, self.base_path, self.protocol, strip_file_scheme=False)
+        dest_path = resolve_storage_path(destination, self.base_path, self.protocol, strip_file_scheme=False)
         self.fs.mv(source_path, dest_path, **kwargs)
 
-    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."""
-        if not PYARROW_INSTALLED:
-            raise MissingDependencyError(package="pyarrow", install_package="pyarrow")
+        ensure_pyarrow()
         import pyarrow.parquet as pq
 
-        resolved_path = self._resolve_path(path)
+        resolved_path = resolve_storage_path(path, self.base_path, self.protocol, strip_file_scheme=False)
         with self.fs.open(resolved_path, mode="rb", **kwargs) as f:
             return pq.read_table(f)
 
-    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."""
-        if not PYARROW_INSTALLED:
-            raise MissingDependencyError(package="pyarrow", install_package="pyarrow")
+        ensure_pyarrow()
         import pyarrow.parquet as pq
 
-        resolved_path = self._resolve_path(path)
+        resolved_path = resolve_storage_path(path, self.base_path, self.protocol, strip_file_scheme=False)
         with self.fs.open(resolved_path, mode="wb") as f:
             pq.write_table(table, f, **kwargs)  # pyright: ignore
 
     def list_objects(self, prefix: str = "", recursive: bool = True, **kwargs: Any) -> list[str]:
         """List objects with optional prefix."""
-        resolved_prefix = self._resolve_path(prefix)
+        resolved_prefix = resolve_storage_path(prefix, self.base_path, self.protocol, strip_file_scheme=False)
         if recursive:
             return sorted(self.fs.find(resolved_prefix, **kwargs))
         return sorted(self.fs.ls(resolved_prefix, detail=False, **kwargs))
 
     def glob(self, pattern: str, **kwargs: Any) -> list[str]:
         """Find objects matching a glob pattern."""
-        resolved_pattern = self._resolve_path(pattern)
+        resolved_pattern = resolve_storage_path(pattern, self.base_path, self.protocol, strip_file_scheme=False)
         return sorted(self.fs.glob(resolved_pattern, **kwargs))  # pyright: ignore
 
-    def is_object(self, path: Union[str, Path]) -> bool:
+    def is_object(self, path: str | Path) -> bool:
         """Check if path points to an object."""
-        resolved_path = self._resolve_path(path)
+        resolved_path = resolve_storage_path(path, self.base_path, self.protocol, strip_file_scheme=False)
         return self.fs.exists(resolved_path) and not self.fs.isdir(resolved_path)
 
-    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)."""
-        resolved_path = self._resolve_path(path)
+        resolved_path = resolve_storage_path(path, self.base_path, self.protocol, strip_file_scheme=False)
         return self.fs.isdir(resolved_path)  # type: ignore[no-any-return]
 
-    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."""
+        resolved_path = resolve_storage_path(path, self.base_path, self.protocol, strip_file_scheme=False)
         try:
-            resolved_path = self._resolve_path(path)
             info = self.fs.info(resolved_path, **kwargs)
+        except FileNotFoundError:
+            return {"path": resolved_path, "exists": False}
+        else:
             if isinstance(info, dict):
                 return {
                     "path": resolved_path,
@@ -229,23 +261,20 @@ class FSSpecBackend:
                     "last_modified": info.get("mtime"),
                     "type": info.get("type", "file"),
                 }
-
-        except FileNotFoundError:
-            return {"path": self._resolve_path(path), "exists": False}
-        return {
-            "path": resolved_path,
-            "exists": True,
-            "size": info.size,
-            "last_modified": info.mtime,
-            "type": info.type,
-        }
+            return {
+                "path": resolved_path,
+                "exists": True,
+                "size": info.size,
+                "last_modified": info.mtime,
+                "type": info.type,
+            }
 
     def sign(self, path: str, expires_in: int = 3600, for_upload: bool = False) -> str:
         """Generate a signed URL for the file."""
-        resolved_path = self._resolve_path(path)
+        resolved_path = resolve_storage_path(path, self.base_path, self.protocol, strip_file_scheme=False)
         return f"{self._fs_uri}{resolved_path}"
 
-    def _stream_file_batches(self, obj_path: Union[str, Path]) -> "Iterator[ArrowRecordBatch]":
+    def _stream_file_batches(self, obj_path: str | Path) -> "Iterator[ArrowRecordBatch]":
         import pyarrow.parquet as pq
 
         with self.fs.open(obj_path, mode="rb") as f:
@@ -253,17 +282,16 @@ class FSSpecBackend:
             yield from parquet_file.iter_batches()
 
     def stream_arrow(self, pattern: str, **kwargs: Any) -> "Iterator[ArrowRecordBatch]":
-        self._ensure_fsspec()
-        self._ensure_pyarrow()
+        ensure_pyarrow()
 
         for obj_path in self.glob(pattern, **kwargs):
             yield from self._stream_file_batches(obj_path)
 
-    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:
         """Read bytes from storage asynchronously."""
         return await async_(self.read_bytes)(path, **kwargs)
 
-    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:
         """Write bytes to storage asynchronously."""
         return await async_(self.write_bytes)(path, data, **kwargs)
 
@@ -277,15 +305,15 @@ class FSSpecBackend:
         Returns:
             AsyncIterator of Arrow record batches
         """
-        self._ensure_pyarrow()
+        ensure_pyarrow()
 
         return _ArrowStreamer(self, pattern, **kwargs)
 
-    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:
         """Read text from storage asynchronously."""
         return await async_(self.read_text)(path, encoding, **kwargs)
 
-    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:
         """Write text to storage asynchronously."""
         await async_(self.write_text)(path, data, encoding, **kwargs)
 
@@ -293,23 +321,23 @@ class FSSpecBackend:
         """List objects in storage asynchronously."""
         return await async_(self.list_objects)(prefix, recursive, **kwargs)
 
-    async def exists_async(self, path: Union[str, Path], **kwargs: Any) -> bool:
+    async def exists_async(self, path: str | Path, **kwargs: Any) -> bool:
         """Check if object exists in storage asynchronously."""
         return await async_(self.exists)(path, **kwargs)
 
-    async def delete_async(self, path: Union[str, Path], **kwargs: Any) -> None:
+    async def delete_async(self, path: str | Path, **kwargs: Any) -> None:
         """Delete object from storage asynchronously."""
         await async_(self.delete)(path, **kwargs)
 
-    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:
         """Copy object in storage asynchronously."""
         await async_(self.copy)(source, destination, **kwargs)
 
-    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:
         """Move object in storage asynchronously."""
         await async_(self.move)(source, destination, **kwargs)
 
-    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]:
         """Get object metadata from storage asynchronously."""
         return await async_(self.get_metadata)(path, **kwargs)
 
@@ -317,10 +345,10 @@ class FSSpecBackend:
         """Generate a signed URL asynchronously."""
         return await async_(self.sign)(path, expires_in, for_upload)
 
-    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":
         """Read Arrow table from storage asynchronously."""
         return await async_(self.read_arrow)(path, **kwargs)
 
-    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:
         """Write Arrow table to storage asynchronously."""
         await async_(self.write_arrow)(path, table, **kwargs)