"""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, cast
from mypy_extensions import mypyc_attr
from typing_extensions import Self
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()
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
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()
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(None, 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)
def __anext__(self) -> Any:
# Returning a Future avoids mypyc coroutine state machine bugs entirely.
return asyncio.get_running_loop().run_in_executor(None, self._sync_next)
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
def __anext__(self) -> Any:
return self._stream.__anext__()
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()
async def aclose(self) -> None:
if self._closed:
return
self._closed = True
with contextlib.suppress(Exception):
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)
def __anext__(self) -> Any:
return asyncio.get_running_loop().run_in_executor(None, self._sync_read)
@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__ = ()
@abstractmethod
def resolve_uri(self, path: "str | Path") -> str:
"""Resolve a backend-relative path to its unsigned address."""
raise NotImplementedError