"""Object storage backend using obstore.
Implements the ObjectStoreProtocol using obstore for S3, GCS, Azure,
and local file storage.
"""
import fnmatch
import io
import re
from collections.abc import AsyncIterator, Iterator
from datetime import timedelta
from functools import partial
from pathlib import Path, PurePosixPath
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, overload
from urllib.parse import urlparse
from mypy_extensions import mypyc_attr
from typing_extensions import Self
from sqlspec.exceptions import StorageOperationFailedError
from sqlspec.storage._arrow_stream import iter_parquet_row_groups, validate_parquet_stream_options
from sqlspec.storage._paths import is_file_destination, resolve_storage_path
from sqlspec.storage._utils import _log_storage_event, import_pyarrow, import_pyarrow_parquet
from sqlspec.storage.backends.base import AsyncArrowBatchIterator, AsyncObStoreStreamIterator
if TYPE_CHECKING:
from obstore.store import ObjectStore
from sqlspec.typing import ArrowRecordBatch, ArrowTable
from sqlspec.storage.errors import execute_sync_storage_operation
from sqlspec.utils.module_loader import ensure_obstore
from sqlspec.utils.sync_tools import async_
DEFAULT_OPTIONS: Final[dict[str, Any]] = {"connect_timeout": "30s", "request_timeout": "60s"}
_MAX_SIGN_EXPIRES_SECONDS: Final[int] = 604800
_SIGNABLE_PROTOCOLS: Final[frozenset[str]] = frozenset({"s3", "gs", "gcs", "az", "azure"})
__all__ = ("ObStoreBackend",)
class _ObStoreFileProxy:
"""Complete obstore's seekable reader interface for PyArrow."""
__slots__ = ("_closed", "_reader")
def __init__(self, reader: Any) -> None:
self._reader = reader
self._closed = False
@property
def closed(self) -> bool:
return self._closed
def readable(self) -> bool:
return not self._closed
def seekable(self) -> bool:
return not self._closed and bool(self._reader.seekable())
def writable(self) -> bool:
return False
def read(self, size: int = -1) -> bytes:
if size < 0:
return cast("bytes", self._reader.readall())
return cast("bytes", self._reader.read(size))
def readinto(self, buffer: Any) -> int:
data = self.read(len(buffer))
buffer[: len(data)] = data
return len(data)
def seek(self, offset: int, whence: int = 0) -> int:
return cast("int", self._reader.seek(offset, whence))
def tell(self) -> int:
return cast("int", self._reader.tell())
def close(self) -> None:
if not self._closed:
self._closed = True
self._reader.close()
def __enter__(self) -> Self:
return self
def __exit__(self, *_: Any) -> None:
self.close()
@mypyc_attr(allow_interpreted_subclasses=True)
class ObStoreBackend:
"""Object storage backend using obstore.
Implements ObjectStoreProtocol using obstore's Rust-based implementation
for storage operations. Supports AWS S3, Google Cloud Storage, Azure Blob Storage,
local filesystem, and HTTP endpoints.
All synchronous methods use the *_sync suffix for consistency with async methods.
Implementation Details & Invariants:
- LocalStore Paths: For LocalStore, the base_path is already included in the store
root (combined with the URI path; if base_path is absolute, Path division will
use it directly). Hence, we use an empty prefix when resolving paths for LocalStore,
whereas cloud stores use base_path as a prefix.
- Native Streaming: Uses obstore's native streaming yielding Buffer objects, which
are converted to bytes.
- Seekable Streams: PyArrow's ParquetFile reads through obstore's seekable
``open_reader`` interface without draining the object into memory.
- Thread Offloading: Uses async_() with a storage limiter to offload blocking
PyArrow serialization/parsing to a thread pool, preventing event loop blocking.
"""
__slots__ = ("_is_local_store", "_local_store_root", "base_path", "protocol", "store", "store_options", "store_uri")
backend_type: ClassVar[str] = "obstore"
def __init__(self, uri: str, **kwargs: Any) -> None:
"""Initialize obstore backend.
Args:
uri: Storage URI. Supported formats:
- file:///absolute/path - Local filesystem
- s3://bucket/prefix - AWS S3
- gs://bucket/prefix - Google Cloud Storage
- az://container/prefix - Azure Blob Storage
- memory:// - In-memory storage (for testing)
**kwargs: Additional options:
- base_path (str): For local files (file://), this is combined with
the URI path to form the storage root. For example:
uri="file:///data" + base_path="uploads" → /data/uploads
If base_path is absolute, it overrides the URI path (backward compat).
For cloud storage, base_path is used as an object key prefix.
- Other obstore configuration options (timeouts, credentials, etc.)
"""
ensure_obstore()
base_path = kwargs.pop("base_path", "")
self.store_uri = uri
self.base_path = base_path.rstrip("/") if base_path else ""
self.store_options = kwargs
self.store: ObjectStore | Any
self._is_local_store = False
self._local_store_root = ""
self.protocol = uri.split("://", 1)[0] if "://" in uri else "file"
try:
if uri.startswith("memory://"):
from obstore.store import MemoryStore
self.store = MemoryStore()
elif uri.startswith("file://"):
from obstore.store import LocalStore
parsed = urlparse(uri)
path_str = parsed.path or "/"
if parsed.fragment:
path_str = f"{path_str}#{parsed.fragment}"
path_obj = Path(path_str)
if is_file_destination(path_obj):
path_str = str(path_obj.parent)
local_store_root_obj = Path(path_str)
if self.base_path:
local_store_root_obj /= self.base_path
self._is_local_store = True
self._local_store_root = str(local_store_root_obj.resolve())
self.store = LocalStore(self._local_store_root, mkdir=True)
else:
from obstore.store import from_url
self.store = from_url(uri, **kwargs) # pyright: ignore[reportAttributeAccessIssue]
_log_storage_event(
"storage.backend.ready",
backend_type=self.backend_type,
protocol=self.protocol,
operation="init",
mode="sync",
path=uri,
)
except Exception as exc:
msg = f"Failed to initialize obstore backend for {uri}"
raise StorageOperationFailedError(msg) from exc