"""Base class for storage backends."""
# ruff: noqa: RSE102
import asyncio
import builtins
import contextlib
from abc import abstractmethod
from collections.abc import AsyncIterator, Iterator
from typing import TYPE_CHECKING, Any, Literal, Self, cast, overload
from mypy_extensions import mypyc_attr
from sqlspec.utils.sync_tools import get_default_async_executor
if TYPE_CHECKING:
from pathlib import Path
from types import TracebackType
from sqlspec.typing import ArrowRecordBatch, ArrowTable
__all__ = ("AsyncArrowBatchIterator", "AsyncObStoreStreamIterator", "AsyncThreadedBytesIterator", "ObjectStoreBase")
_StopAsyncBase = getattr(builtins, "Stop" + "Async" + "Iteration")
_StopAsync = type("_StopAsync", (_StopAsyncBase,), {})
class _ExhaustedSentinel:
"""Sentinel value to signal iterator exhaustion across thread boundaries.
StopIteration cannot be raised into asyncio Futures, so we use this sentinel
to signal iterator exhaustion from the thread pool back to the async context.
"""
__slots__ = ()
_EXHAUSTED = _ExhaustedSentinel()
class AsyncArrowBatchIterator:
"""Async iterator wrapper for sync Arrow batch iterators."""
__slots__ = ("_closed", "_sync_iter")
def __init__(self, sync_iterator: "Iterator[ArrowRecordBatch]") -> None:
self._sync_iter = sync_iterator
self._closed = False
def __aiter__(self) -> "AsyncArrowBatchIterator":
return self
async def __aenter__(self) -> Self:
return self
async def __aexit__(
self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
) -> None:
await self.aclose()
def __anext__(self) -> Any:
return asyncio.get_running_loop().run_in_executor(get_default_async_executor(), self._sync_next)
async def aclose(self) -> None:
"""Close the underlying generator and its active storage reader."""
if self._closed:
return
self._closed = True
close = getattr(self._sync_iter, "close", None)
if close is not None:
await asyncio.get_running_loop().run_in_executor(get_default_async_executor(), close)
def _sync_next(self) -> "ArrowRecordBatch":
if self._closed:
raise _StopAsync()
result = _next_or_sentinel(self._sync_iter)
if result is _EXHAUSTED:
self._closed = True
raise _StopAsync()
return cast("ArrowRecordBatch", result)
class AsyncObStoreStreamIterator:
"""Async iterator wrapper for obstore streaming."""
__slots__ = ("_chunk_size", "_stream")
def __init__(self, stream: Any, chunk_size: "int | None" = None) -> None:
self._stream = stream
self._chunk_size = chunk_size if chunk_size is not None and chunk_size > 0 else None
def __aiter__(self) -> "AsyncObStoreStreamIterator":
return self
async def __aenter__(self) -> Self:
return self
async def __aexit__(
self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
) -> None:
await self.aclose()
async def __anext__(self) -> bytes:
if self._stream is None:
raise StopAsyncIteration
chunk = await self._stream.__anext__()
return chunk if type(chunk) is bytes else bytes(chunk)
async def aclose(self) -> None:
"""Close the underlying stream if it exposes a close hook."""
stream = self._stream
if stream is None:
return
self._stream = None
aclose: Any = getattr(stream, "aclose", None)
if aclose is not None:
with contextlib.suppress(Exception):
await aclose()
return
close: Any = getattr(stream, "close", None)
if close is not None:
with contextlib.suppress(Exception):
close()
class AsyncThreadedBytesIterator:
"""Async iterator that reads from a synchronous file-like object in a thread pool."""
__slots__ = ("_chunk_size", "_closed", "_file_obj")
def __init__(self, file_obj: Any, chunk_size: int = 65536) -> None:
self._file_obj = file_obj
self._chunk_size = chunk_size
self._closed = False
def __aiter__(self) -> "AsyncThreadedBytesIterator":
return self
async def __aenter__(self) -> Self:
return self
async def __aexit__(
self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
) -> None:
await self.aclose()
def __anext__(self) -> Any:
return asyncio.get_running_loop().run_in_executor(get_default_async_executor(), self._sync_read)
async def aclose(self) -> None:
if self._closed:
return
self._closed = True
with contextlib.suppress(Exception):
await asyncio.get_running_loop().run_in_executor(get_default_async_executor(), self._file_obj.close)
def _sync_read(self) -> bytes:
if self._closed:
raise _StopAsync()
result = _read_chunk_or_sentinel(self._file_obj, self._chunk_size)
if result is _EXHAUSTED:
self._closed = True
with contextlib.suppress(Exception):
self._file_obj.close()
raise _StopAsync()
return cast("bytes", result)
[docs]
@mypyc_attr(allow_interpreted_subclasses=True)
class ObjectStoreBase:
"""Base class for storage backends.
All synchronous methods follow the *_sync naming convention for consistency
with their async counterparts.
"""
__slots__ = ()
@property
@abstractmethod
def supports_signing(self) -> bool:
"""Whether this backend supports URL signing."""
raise NotImplementedError
[docs]
@abstractmethod
def resolve_uri(self, path: "str | Path") -> str:
"""Resolve a backend-relative path to its unsigned address."""
raise NotImplementedError
[docs]
@abstractmethod
def read_bytes_sync(self, path: "str | Path", **kwargs: Any) -> bytes:
"""Read bytes from storage synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def write_bytes_sync(self, path: "str | Path", data: bytes, **kwargs: Any) -> None:
"""Write bytes to storage synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def stream_read_sync(self, path: "str | Path", chunk_size: "int | None" = None, **kwargs: Any) -> Iterator[bytes]:
"""Stream bytes from storage synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def read_text_sync(self, path: "str | Path", encoding: str = "utf-8", **kwargs: Any) -> str:
"""Read text from storage synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def write_text_sync(self, path: "str | Path", data: str, encoding: str = "utf-8", **kwargs: Any) -> None:
"""Write text to storage synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def list_objects_sync(self, prefix: str = "", recursive: bool = True, **kwargs: Any) -> "list[str]":
"""List objects in storage synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def exists_sync(self, path: "str | Path", **kwargs: Any) -> bool:
"""Check if object exists in storage synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def delete_sync(self, path: "str | Path", **kwargs: Any) -> None:
"""Delete object from storage synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def copy_sync(self, source: "str | Path", destination: "str | Path", **kwargs: Any) -> None:
"""Copy object within storage synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def move_sync(self, source: "str | Path", destination: "str | Path", **kwargs: Any) -> None:
"""Move object within storage synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def glob_sync(self, pattern: str, **kwargs: Any) -> "list[str]":
"""Find objects matching pattern synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def is_object_sync(self, path: "str | Path") -> bool:
"""Check if path points to an object synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def is_path_sync(self, path: "str | Path") -> bool:
"""Check if path points to a directory synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def read_arrow_sync(self, path: "str | Path", **kwargs: Any) -> "ArrowTable":
"""Read Arrow table from storage synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def write_arrow_sync(self, path: "str | Path", table: "ArrowTable", **kwargs: Any) -> None:
"""Write Arrow table to storage synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def stream_arrow_sync(
self, pattern: str, *, file_format: Literal["parquet"] = "parquet", batch_size: int = 65_536, **kwargs: Any
) -> "Iterator[ArrowRecordBatch]":
"""Stream Arrow record batches from storage synchronously."""
raise NotImplementedError
@overload
def sign_sync(self, paths: str, expires_in: int = 3600, for_upload: bool = False) -> str: ...
@overload
def sign_sync(self, paths: "list[str]", expires_in: int = 3600, for_upload: bool = False) -> "list[str]": ...
[docs]
@abstractmethod
def sign_sync(
self, paths: "str | list[str]", expires_in: int = 3600, for_upload: bool = False
) -> "str | list[str]":
"""Generate signed URL(s) synchronously."""
raise NotImplementedError
[docs]
@abstractmethod
async def read_bytes_async(self, path: "str | Path", **kwargs: Any) -> bytes:
"""Read bytes from storage asynchronously."""
raise NotImplementedError
[docs]
@abstractmethod
async def write_bytes_async(self, path: "str | Path", data: bytes, **kwargs: Any) -> None:
"""Write bytes to storage asynchronously."""
raise NotImplementedError
[docs]
@abstractmethod
async def read_text_async(self, path: "str | Path", encoding: str = "utf-8", **kwargs: Any) -> str:
"""Read text from storage asynchronously."""
raise NotImplementedError
[docs]
@abstractmethod
async def write_text_async(self, path: "str | Path", data: str, encoding: str = "utf-8", **kwargs: Any) -> None:
"""Write text to storage asynchronously."""
raise NotImplementedError
[docs]
@abstractmethod
async def stream_read_async(
self, path: "str | Path", chunk_size: "int | None" = None, **kwargs: Any
) -> AsyncIterator[bytes]:
"""Stream bytes from storage asynchronously."""
raise NotImplementedError
[docs]
@abstractmethod
async def list_objects_async(self, prefix: str = "", recursive: bool = True, **kwargs: Any) -> "list[str]":
"""List objects in storage asynchronously."""
raise NotImplementedError
[docs]
@abstractmethod
async def exists_async(self, path: "str | Path", **kwargs: Any) -> bool:
"""Check if object exists in storage asynchronously."""
raise NotImplementedError
[docs]
@abstractmethod
async def delete_async(self, path: "str | Path", **kwargs: Any) -> None:
"""Delete object from storage asynchronously."""
raise NotImplementedError
[docs]
@abstractmethod
async def copy_async(self, source: "str | Path", destination: "str | Path", **kwargs: Any) -> None:
"""Copy object within storage asynchronously."""
raise NotImplementedError
[docs]
@abstractmethod
async def move_async(self, source: "str | Path", destination: "str | Path", **kwargs: Any) -> None:
"""Move object within storage asynchronously."""
raise NotImplementedError
[docs]
@abstractmethod
async def read_arrow_async(self, path: "str | Path", **kwargs: Any) -> "ArrowTable":
"""Read Arrow table from storage asynchronously."""
raise NotImplementedError
[docs]
@abstractmethod
async def write_arrow_async(self, path: "str | Path", table: "ArrowTable", **kwargs: Any) -> None:
"""Write Arrow table to storage asynchronously."""
raise NotImplementedError
[docs]
@abstractmethod
def stream_arrow_async(
self, pattern: str, *, file_format: Literal["parquet"] = "parquet", batch_size: int = 65_536, **kwargs: Any
) -> "AsyncIterator[ArrowRecordBatch]":
"""Stream Arrow record batches from storage asynchronously.
Returns AsyncIterator directly (not a coroutine).
"""
raise NotImplementedError
@overload
async def sign_async(self, paths: str, expires_in: int = 3600, for_upload: bool = False) -> str: ...
@overload
async def sign_async(self, paths: "list[str]", expires_in: int = 3600, for_upload: bool = False) -> "list[str]": ...
[docs]
@abstractmethod
async def sign_async(
self, paths: "str | list[str]", expires_in: int = 3600, for_upload: bool = False
) -> "str | list[str]":
"""Generate signed URL(s) asynchronously."""
raise NotImplementedError
def _next_or_sentinel(iterator: "Iterator[Any]") -> "Any":
"""Get next item or return sentinel if exhausted."""
try:
return next(iterator)
except StopIteration:
return _EXHAUSTED
def _read_chunk_or_sentinel(file_obj: Any, chunk_size: int) -> Any:
"""Read a chunk from a file-like object or return sentinel if exhausted."""
try:
chunk = file_obj.read(chunk_size)
if not chunk:
return _EXHAUSTED
except EOFError:
return _EXHAUSTED
return chunk
