Skip to main content

自定义后端

如何实现自定义存储后端。自定义后端需实现 StorageBackend 接口,可被队列系统内部使用。

StorageBackend 协议

实现以下方法(定义在 persisting.queue.backend):

from typing import Any, AsyncIterator

class MyBackend:
def __init__(self, bucket_id: int, storage_path: str, batch_size: int = 100, **kwargs):
...

async def put(self, record: dict[str, Any]) -> None: ...
async def put_batch(self, records: list[dict[str, Any]]) -> None: ...
async def get(self, limit: int, offset: int) -> list[dict[str, Any]]: ...
async def get_stream(self, limit: int, offset: int, wait: bool = False, timeout: float | None = None) -> AsyncIterator[list[dict[str, Any]]]: ...
async def flush(self) -> None: ...
async def stats(self) -> dict[str, Any]: ...
def total_count(self) -> int: ...

构造函数必须接受 bucket_idstorage_pathbatch_size,以及 **kwargs 以兼容未来新增参数。

最小示例

import asyncio
from typing import Any, AsyncIterator

class SimpleBackend:
def __init__(self, bucket_id: int, storage_path: str, batch_size: int = 100, **kwargs):
self.bucket_id = bucket_id
self._records: list[dict[str, Any]] = []
self._condition = asyncio.Condition()

def total_count(self) -> int:
return len(self._records)

async def put(self, record: dict[str, Any]) -> None:
async with self._condition:
self._records.append(record)
self._condition.notify_all()

async def put_batch(self, records: list[dict[str, Any]]) -> None:
async with self._condition:
self._records.extend(records)
self._condition.notify_all()

async def get(self, limit: int, offset: int) -> list[dict[str, Any]]:
return self._records[offset : offset + limit]

async def get_stream(
self, limit: int, offset: int, wait: bool = False, timeout: float | None = None,
) -> AsyncIterator[list[dict[str, Any]]]:
pos = offset
remaining = limit
while remaining > 0:
async with self._condition:
if pos >= len(self._records):
if not wait:
return
try:
coro = self._condition.wait()
if timeout:
await asyncio.wait_for(coro, timeout=timeout)
else:
await coro
continue
except asyncio.TimeoutError:
return
batch = self._records[pos : pos + min(remaining, 100)]
if batch:
yield batch
pos += len(batch)
remaining -= len(batch)
elif not wait:
break

async def flush(self) -> None:
pass # 内存存储无需 flush

async def stats(self) -> dict[str, Any]:
return {"bucket_id": self.bucket_id, "total_count": len(self._records)}

实现此接口的自定义后端可在接受后端实例的地方内部使用。公开 API(Queue)使用基于 Lance 的存储;自定义后端适用于高级场景和框架集成。

要点

  1. 使用 asyncio.Condition 保证并发安全和阻塞读取。
  2. 构造函数接受 **kwargs 以兼容未来新增的选项。
  3. 写入后调用 notify_all()get_stream(wait=True) 的读取端能被唤醒。
  4. get_stream 返回批次list[dict]),不是单条记录。

下一步