add:数据库添加批量操作支持,附测试代码

This commit is contained in:
Lishiling
2026-03-13 11:02:31 +08:00
parent da2c252b55
commit 945bc3bc34
7 changed files with 407 additions and 7 deletions

8
.gitignore vendored
View File

@@ -7,6 +7,7 @@ __pycache__/
*.pyd
*.so
.pytest_cache/
pytest-cache-files-*/
.mypy_cache/
.ruff_cache/
.coverage
@@ -22,6 +23,13 @@ wheels/
.eggs/
pip-wheel-metadata/
#
fork-docs/
tmp/
openspec/
scripts/
docs/zh/reference
# Virtual environments
.venv/
venv/

View File

@@ -18,14 +18,11 @@
- 数据永久存储,除非用户显式删除
- 值类型支持任意 JSON 数据
- 支持前缀查询键列表
TODO:
- 缺少批量操作支持 (set_many, get_many)
- 缺少数据变更事件通知
"""
from __future__ import annotations
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import Any
from ._proxy import CapabilityProxy
@@ -108,3 +105,67 @@ class DBClient:
if not isinstance(keys, (list, tuple)):
return []
return [str(item) for item in keys]
async def get_many(self, keys: Sequence[str]) -> dict[str, Any | None]:
"""批量获取多个键的值。
Args:
keys: 要读取的键列表
Returns:
一个 dictkey 为键名value 为对应值(不存在则为 None
示例:
values = await ctx.db.get_many(["user:1", "user:2"])
if values["user:1"] is None:
print("user:1 missing")
"""
output = await self._proxy.call("db.get_many", {"keys": list(keys)})
items = output.get("items")
if not isinstance(items, (list, tuple)):
return {}
result: dict[str, Any | None] = {}
for item in items:
if not isinstance(item, dict):
continue
key = item.get("key")
if not isinstance(key, str):
continue
result[key] = item.get("value")
return result
async def set_many(
self, items: Mapping[str, Any] | Sequence[tuple[str, Any]]
) -> None:
"""批量写入多个键值对。
Args:
items: 键值对集合dict 或二元组序列)
示例:
await ctx.db.set_many({"user:1": {"name": "a"}, "user:2": {"name": "b"}})
"""
if isinstance(items, Mapping):
pairs = list(items.items())
else:
pairs = list(items)
payload_items: list[dict[str, Any]] = [
{"key": str(key), "value": value} for key, value in pairs
]
await self._proxy.call("db.set_many", {"items": payload_items})
def watch(self, prefix: str | None = None) -> AsyncIterator[dict[str, Any]]:
"""订阅 KV 变更事件(流式)。
Args:
prefix: 键前缀过滤None 表示订阅所有键
Yields:
变更事件 dict{"op": "set"|"delete", "key": str, "value": Any|None}
示例:
async for event in ctx.db.watch("user:"):
print(event["op"], event["key"])
"""
return self._proxy.stream("db.watch", {"prefix": prefix})

View File

@@ -127,6 +127,37 @@ DB_LIST_OUTPUT_SCHEMA = _object_schema(
required=("keys",),
keys={"type": "array", "items": {"type": "string"}},
)
DB_GET_MANY_INPUT_SCHEMA = _object_schema(
required=("keys",),
keys={"type": "array", "items": {"type": "string"}},
)
DB_GET_MANY_OUTPUT_SCHEMA = _object_schema(
required=("items",),
items={
"type": "array",
"items": _object_schema(
required=("key", "value"),
key={"type": "string"},
value=_nullable({}),
),
},
)
DB_SET_MANY_INPUT_SCHEMA = _object_schema(
required=("items",),
items={
"type": "array",
"items": _object_schema(
required=("key", "value"),
key={"type": "string"},
value={},
),
},
)
DB_SET_MANY_OUTPUT_SCHEMA = _object_schema()
DB_WATCH_INPUT_SCHEMA = _object_schema(
prefix=_nullable({"type": "string"}),
)
DB_WATCH_OUTPUT_SCHEMA = _object_schema()
SESSION_REF_SCHEMA = _object_schema(
required=("conversation_id",),
conversation_id={"type": "string"},
@@ -218,6 +249,18 @@ BUILTIN_CAPABILITY_SCHEMAS: dict[str, dict[str, JSONSchema]] = {
"input": DB_LIST_INPUT_SCHEMA,
"output": DB_LIST_OUTPUT_SCHEMA,
},
"db.get_many": {
"input": DB_GET_MANY_INPUT_SCHEMA,
"output": DB_GET_MANY_OUTPUT_SCHEMA,
},
"db.set_many": {
"input": DB_SET_MANY_INPUT_SCHEMA,
"output": DB_SET_MANY_OUTPUT_SCHEMA,
},
"db.watch": {
"input": DB_WATCH_INPUT_SCHEMA,
"output": DB_WATCH_OUTPUT_SCHEMA,
},
"platform.send": {
"input": PLATFORM_SEND_INPUT_SCHEMA,
"output": PLATFORM_SEND_OUTPUT_SCHEMA,
@@ -491,10 +534,16 @@ __all__ = [
"DB_DELETE_OUTPUT_SCHEMA",
"DB_GET_INPUT_SCHEMA",
"DB_GET_OUTPUT_SCHEMA",
"DB_GET_MANY_INPUT_SCHEMA",
"DB_GET_MANY_OUTPUT_SCHEMA",
"DB_LIST_INPUT_SCHEMA",
"DB_LIST_OUTPUT_SCHEMA",
"DB_SET_INPUT_SCHEMA",
"DB_SET_OUTPUT_SCHEMA",
"DB_SET_MANY_INPUT_SCHEMA",
"DB_SET_MANY_OUTPUT_SCHEMA",
"DB_WATCH_INPUT_SCHEMA",
"DB_WATCH_OUTPUT_SCHEMA",
"EventTrigger",
"HandlerDescriptor",
"JSONSchema",

View File

@@ -109,6 +109,7 @@ CAPABILITY_NAME_PATTERN = re.compile(r"^[a-z][a-z0-9_]*\.[a-z][a-z0-9_]*$")
class StreamExecution:
iterator: AsyncIterator[dict[str, Any]]
finalize: FinalizeHandler
collect_chunks: bool = True
StreamHandler = Callable[
@@ -134,8 +135,18 @@ class CapabilityRouter:
self.db_store: dict[str, Any] = {}
self.memory_store: dict[str, dict[str, Any]] = {}
self.sent_messages: list[dict[str, Any]] = []
self._db_watch_subscriptions: dict[
str, tuple[str | None, asyncio.Queue[dict[str, Any]]]
] = {}
self._register_builtin_capabilities()
def _emit_db_change(self, *, op: str, key: str, value: Any | None) -> None:
event = {"op": op, "key": key, "value": value}
for prefix, queue in list(self._db_watch_subscriptions.values()):
if prefix is not None and not key.startswith(prefix):
continue
queue.put_nowait(event)
def descriptors(self) -> list[CapabilityDescriptor]:
return [
entry.descriptor for entry in self._registrations.values() if entry.exposed
@@ -227,6 +238,7 @@ class CapabilityRouter:
return StreamExecution(
iterator=execution.iterator,
finalize=validated_finalize,
collect_chunks=execution.collect_chunks,
)
def _register_builtin_capabilities(self) -> None:
@@ -329,13 +341,17 @@ class CapabilityRouter:
_request_id: str, payload: dict[str, Any], _token
) -> dict[str, Any]:
key = str(payload.get("key", ""))
self.db_store[key] = payload.get("value")
value = payload.get("value")
self.db_store[key] = value
self._emit_db_change(op="set", key=key, value=value)
return {}
async def db_delete(
_request_id: str, payload: dict[str, Any], _token
) -> dict[str, Any]:
self.db_store.pop(str(payload.get("key", "")), None)
key = str(payload.get("key", ""))
self.db_store.pop(key, None)
self._emit_db_change(op="delete", key=key, value=None)
return {}
async def db_list(
@@ -347,6 +363,63 @@ class CapabilityRouter:
keys = [item for item in keys if item.startswith(prefix)]
return {"keys": keys}
async def db_get_many(
_request_id: str, payload: dict[str, Any], _token
) -> dict[str, Any]:
keys_payload = payload.get("keys")
if not isinstance(keys_payload, (list, tuple)):
raise AstrBotError.invalid_input("db.get_many 的 keys 必须是数组")
keys = [str(item) for item in keys_payload]
items = [{"key": key, "value": self.db_store.get(key)} for key in keys]
return {"items": items}
async def db_set_many(
_request_id: str, payload: dict[str, Any], _token
) -> dict[str, Any]:
items_payload = payload.get("items")
if not isinstance(items_payload, (list, tuple)):
raise AstrBotError.invalid_input("db.set_many 的 items 必须是数组")
for entry in items_payload:
if not isinstance(entry, dict):
raise AstrBotError.invalid_input(
"db.set_many 的 items 必须是 object 数组"
)
key = str(entry.get("key", ""))
value = entry.get("value")
self.db_store[key] = value
self._emit_db_change(op="set", key=key, value=value)
return {}
async def db_watch(
request_id: str, payload: dict[str, Any], _token
) -> StreamExecution:
prefix = payload.get("prefix")
prefix_value: str | None
if isinstance(prefix, str):
prefix_value = prefix
elif prefix is None:
prefix_value = None
else:
raise AstrBotError.invalid_input(
"db.watch 的 prefix 必须是 string 或 null"
)
queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
self._db_watch_subscriptions[request_id] = (prefix_value, queue)
async def iterator() -> AsyncIterator[dict[str, Any]]:
try:
while True:
yield await queue.get()
finally:
self._db_watch_subscriptions.pop(request_id, None)
return StreamExecution(
iterator=iterator(),
finalize=lambda _chunks: {},
collect_chunks=False,
)
async def platform_send(
_request_id: str, payload: dict[str, Any], _token
) -> dict[str, Any]:
@@ -460,6 +533,23 @@ class CapabilityRouter:
builtin_descriptor("db.list", "列出 KV"),
call_handler=db_list,
)
self.register(
builtin_descriptor("db.get_many", "批量读取 KV"),
call_handler=db_get_many,
)
self.register(
builtin_descriptor("db.set_many", "批量写入 KV"),
call_handler=db_set_many,
)
self.register(
builtin_descriptor(
"db.watch",
"订阅 KV 变更",
supports_stream=True,
cancelable=True,
),
stream_handler=db_watch,
)
self.register(
builtin_descriptor("platform.send", "发送消息"),
call_handler=platform_send,

View File

@@ -507,9 +507,11 @@ class Peer:
"stream=true 必须返回 StreamExecution"
)
await self._send(EventMessage(id=message.id, phase="started"))
collect_chunks = execution.collect_chunks
chunks: list[dict[str, Any]] = []
async for chunk in execution.iterator:
chunks.append(chunk)
if collect_chunks:
chunks.append(chunk)
await self._send(
EventMessage(id=message.id, phase="delta", data=chunk)
)

View File

@@ -164,6 +164,9 @@ class TestCapabilityRouterInit:
assert "db.set" in capability_names
assert "db.delete" in capability_names
assert "db.list" in capability_names
assert "db.get_many" in capability_names
assert "db.set_many" in capability_names
assert "db.watch" in capability_names
# Platform capabilities
assert "platform.send" in capability_names
@@ -182,6 +185,14 @@ class TestCapabilityRouterInit:
assert descriptors[name].input_schema == schema["input"]
assert descriptors[name].output_schema == schema["output"]
def test_db_watch_descriptor_supports_stream(self):
router = CapabilityRouter()
descriptors = {
descriptor.name: descriptor for descriptor in router.descriptors()
}
assert descriptors["db.watch"].supports_stream is True
class TestCapabilityRouterRegister:
"""Tests for CapabilityRouter.register method."""
@@ -350,6 +361,85 @@ class TestCapabilityRouterExecute:
request_id="req-1",
)
class TestCapabilityRouterDBWatch:
"""Router-level tests for db.watch behavior."""
@pytest.mark.asyncio
async def test_db_watch_receives_set_and_delete_events(self):
router = CapabilityRouter()
token = CancelToken()
execution = await router.execute(
"db.watch",
{"prefix": None},
stream=True,
cancel_token=token,
request_id="watch-1",
)
assert isinstance(execution, StreamExecution)
assert execution.collect_chunks is False
await router.execute(
"db.set",
{"key": "a", "value": 1},
stream=False,
cancel_token=token,
request_id="set-1",
)
await router.execute(
"db.delete",
{"key": "a"},
stream=False,
cancel_token=token,
request_id="del-1",
)
event1 = await anext(execution.iterator)
event2 = await anext(execution.iterator)
assert event1 == {"op": "set", "key": "a", "value": 1}
assert event2 == {"op": "delete", "key": "a", "value": None}
close = getattr(execution.iterator, "aclose", None)
if close is not None:
await close()
@pytest.mark.asyncio
async def test_db_watch_prefix_filters_events(self):
router = CapabilityRouter()
token = CancelToken()
execution = await router.execute(
"db.watch",
{"prefix": "user:"},
stream=True,
cancel_token=token,
request_id="watch-2",
)
assert isinstance(execution, StreamExecution)
await router.execute(
"db.set",
{"key": "sys:1", "value": 1},
stream=False,
cancel_token=token,
request_id="set-2",
)
await router.execute(
"db.set",
{"key": "user:1", "value": {"ok": True}},
stream=False,
cancel_token=token,
request_id="set-3",
)
event = await anext(execution.iterator)
assert event == {"op": "set", "key": "user:1", "value": {"ok": True}}
close = getattr(execution.iterator, "aclose", None)
if close is not None:
await close()
@pytest.mark.asyncio
async def test_execute_missing_capability_raises(self):
"""execute should raise for unknown capability."""

View File

@@ -237,3 +237,103 @@ class TestDBClientList:
proxy.call.assert_called_once_with("db.list", {"prefix": None})
assert result == []
class TestDBClientGetMany:
"""Tests for DBClient.get_many() method."""
@pytest.mark.asyncio
async def test_get_many_returns_mapping(self):
proxy = MagicMock(spec=CapabilityProxy)
proxy.call = AsyncMock(
return_value={
"items": [
{"key": "a", "value": 1},
{"key": "b", "value": None},
]
}
)
client = DBClient(proxy)
result = await client.get_many(["a", "b"])
proxy.call.assert_called_once_with("db.get_many", {"keys": ["a", "b"]})
assert result == {"a": 1, "b": None}
@pytest.mark.asyncio
async def test_get_many_returns_empty_dict_for_malformed_items(self):
proxy = MagicMock(spec=CapabilityProxy)
proxy.call = AsyncMock(return_value={"items": "not-a-list"})
client = DBClient(proxy)
result = await client.get_many(["a"])
assert result == {}
class TestDBClientSetMany:
"""Tests for DBClient.set_many() method."""
@pytest.mark.asyncio
async def test_set_many_accepts_mapping(self):
proxy = MagicMock(spec=CapabilityProxy)
proxy.call = AsyncMock(return_value={})
client = DBClient(proxy)
await client.set_many({"a": 1, "b": 2})
proxy.call.assert_called_once_with(
"db.set_many",
{"items": [{"key": "a", "value": 1}, {"key": "b", "value": 2}]},
)
@pytest.mark.asyncio
async def test_set_many_accepts_sequence_pairs(self):
proxy = MagicMock(spec=CapabilityProxy)
proxy.call = AsyncMock(return_value={})
client = DBClient(proxy)
await client.set_many([("a", True), ("b", {"x": 1})])
proxy.call.assert_called_once_with(
"db.set_many",
{"items": [{"key": "a", "value": True}, {"key": "b", "value": {"x": 1}}]},
)
class TestDBClientWatch:
"""Tests for DBClient.watch() method."""
@pytest.mark.asyncio
async def test_watch_calls_proxy_stream_and_yields_events(self):
async def gen():
yield {"op": "set", "key": "a", "value": 1}
yield {"op": "delete", "key": "a", "value": None}
proxy = MagicMock(spec=CapabilityProxy)
proxy.stream = MagicMock(return_value=gen())
client = DBClient(proxy)
iterator = client.watch()
proxy.stream.assert_called_once_with("db.watch", {"prefix": None})
events = [event async for event in iterator]
assert events == [
{"op": "set", "key": "a", "value": 1},
{"op": "delete", "key": "a", "value": None},
]
@pytest.mark.asyncio
async def test_watch_with_prefix(self):
async def gen():
yield {"op": "set", "key": "user:1", "value": {"ok": True}}
proxy = MagicMock(spec=CapabilityProxy)
proxy.stream = MagicMock(return_value=gen())
client = DBClient(proxy)
iterator = client.watch(prefix="user:")
proxy.stream.assert_called_once_with("db.watch", {"prefix": "user:"})
events = [event async for event in iterator]
assert events == [{"op": "set", "key": "user:1", "value": {"ok": True}}]