Source code for sqlspec.storage.backends.fsspec

# pyright: reportPrivateUsage=false
import asyncio
from collections.abc import AsyncIterator, Iterator
from functools import partial
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast, overload
from urllib.parse import urlparse

from mypy_extensions import mypyc_attr

from sqlspec.storage._arrow_stream import iter_parquet_row_groups, validate_parquet_stream_options
from sqlspec.storage._paths import resolve_storage_path
from sqlspec.storage._utils import _log_storage_event, import_pyarrow_parquet
from sqlspec.storage.backends.base import AsyncArrowBatchIterator, AsyncThreadedBytesIterator
from sqlspec.storage.errors import execute_sync_storage_operation
from sqlspec.utils.module_loader import ensure_fsspec
from sqlspec.utils.sync_tools import async_

if TYPE_CHECKING:
    from sqlspec.typing import ArrowRecordBatch, ArrowTable

__all__ = ("FSSpecBackend",)

_OBJECT_STORE_PROTOCOLS = {"s3", "gs", "az", "gcs"}




@mypyc_attr(allow_interpreted_subclasses=True)
class FSSpecBackend:
    """Storage backend using fsspec.

    Implements ObjectStoreProtocol using fsspec for various protocols
    including HTTP, HTTPS, FTP, and cloud storage services.

    All synchronous methods use the *_sync suffix for consistency with async methods.
    """

    __slots__ = ("_fs_uri", "base_path", "fs", "protocol")

    backend_type: ClassVar[str] = "fsspec"



    def __init__(self, uri: str, **kwargs: Any) -> None:
        """Initialize the fsspec-backed storage backend.

        Args:
            uri: Filesystem URI (protocol://path).
            **kwargs: Additional fsspec configuration options, including an optional base_path.

            For cloud URIs (S3/GS/Azure) and file:// URIs, we derive a default base_path from the
                URI path when no explicit base_path is provided. When both URI and base_path are provided,
                they are combined (base_path is appended to URI-derived path).
        """
        ensure_fsspec()
        import fsspec

        explicit_base_path = kwargs.pop("base_path", "")

        if "://" in uri:
            self.protocol = uri.split("://", maxsplit=1)[0]
            self._fs_uri = uri

            if self.protocol in _OBJECT_STORE_PROTOCOLS:
                parsed = urlparse(uri)
                if parsed.netloc:
                    uri_base_path = parsed.netloc
                    if parsed.path and parsed.path != "/":
                        uri_base_path = f"{uri_base_path}{parsed.path}"
                    # Combine URI path with explicit base_path if both provided
                    if explicit_base_path:
                        uri_base_path = f"{uri_base_path.rstrip('/')}/{explicit_base_path.lstrip('/')}"
                    explicit_base_path = uri_base_path
            elif self.protocol == "file":
                parsed = urlparse(uri)
                if parsed.path and parsed.path != "/":
                    # For file protocol, keep the path as-is (preserve leading slash for absolute paths)
                    uri_base_path = parsed.path
                    # Combine URI path with explicit base_path if both provided
                    if explicit_base_path:
                        uri_base_path = f"{uri_base_path.rstrip('/')}/{explicit_base_path.lstrip('/')}"
                    explicit_base_path = uri_base_path
        else:
            self.protocol = uri
            self._fs_uri = f"{uri}://"

        self.base_path = explicit_base_path.rstrip("/") if explicit_base_path else ""

        self.fs = fsspec.filesystem(self.protocol, **kwargs)
        _log_storage_event(
            "storage.backend.ready",
            backend_type=self.backend_type,
            protocol=self.protocol,
            operation="init",
            path=self._fs_uri,
        )

        super().__init__()

@classmethod def from_config(cls, config: "dict[str, Any]") -> "FSSpecBackend": protocol = config["protocol"] fs_config = config.get("fs_config", {}) base_path = config.get("base_path", "") uri = f"{protocol}://" kwargs = dict(fs_config) if base_path: kwargs["base_path"] = base_path return cls(uri=uri, **kwargs) @property def base_uri(self) -> str: return self._fs_uri

def resolve_uri(self, path: str | Path) -> str: """Resolve a backend-relative path to an unsigned address. Args: path: The same backend-relative path accepted by read and write methods. Returns: An absolute filesystem path for ``file`` or a protocol-qualified URI for other filesystems. The target does not need to exist. """ resolved_path = self._resolve_path(path) if self.protocol == "file": return str(Path(resolved_path).resolve()) return str(self.fs.unstrip_protocol(resolved_path))