144 lines
5.7 KiB
Python
144 lines
5.7 KiB
Python
"""DB 持久化事件存储：写 events 表 + 内存 pub/sub 供 SSE 实时推送。
|
||
|
||
混合架构：DB 是事实源（持久化、可回放、重启不丢），内存订阅表提供实时推送
|
||
（DB 无推送能力）。实现 ``EventStore`` Protocol，与 ``InMemoryEventStore`` 可互换，
|
||
调用方无感知（见 events.py 的 Protocol 定义）。
|
||
|
||
sequence 分配：run 内单调递增。``append`` 时若未指定则查 ``MAX(sequence)+1``，
|
||
并发下撞 ``uq_events_run_sequence`` 时重新分配并重试（沙箱场景并发不高，重试足够）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import fnmatch
|
||
import json
|
||
from datetime import datetime
|
||
|
||
from sqlalchemy import func, select
|
||
from sqlalchemy.exc import IntegrityError
|
||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||
|
||
from app.models.runtime import Event as EventRow
|
||
from app.modules.runtime.events import Event
|
||
|
||
# 并发 sequence 冲突时的重试上限
|
||
_MAX_SEQUENCE_RETRIES = 3
|
||
|
||
|
||
class DbEventStore:
|
||
"""DB 持久化 + 内存 pub/sub 的事件存储。"""
|
||
|
||
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
||
self._session_factory = session_factory
|
||
self._subscribers: dict[str, list[asyncio.Queue[Event]]] = {}
|
||
self._pattern_subscribers: dict[str, list[asyncio.Queue[Event]]] = {}
|
||
|
||
async def append(self, event: Event, *, session: AsyncSession | None = None) -> Event:
|
||
# run 内事件分配单调 sequence；run 外事件（run_id=None）sequence 无意义，用 0 占位
|
||
if event.sequence is None and event.run_id:
|
||
event = await self._assign_sequence(event)
|
||
|
||
for attempt in range(_MAX_SEQUENCE_RETRIES + 1):
|
||
row = _event_to_row(event)
|
||
try:
|
||
if session is not None:
|
||
session.add(row)
|
||
# 不提交——调用方管理事务
|
||
else:
|
||
async with self._session_factory() as s:
|
||
s.add(row)
|
||
await s.commit()
|
||
break
|
||
except IntegrityError:
|
||
if attempt < _MAX_SEQUENCE_RETRIES and event.run_id:
|
||
event = await self._assign_sequence(event)
|
||
continue
|
||
raise
|
||
|
||
# 实时通知订阅者（内存 pub/sub，供 SSE 推送）
|
||
for queue in self._subscribers.get(event.run_id or "", []):
|
||
queue.put_nowait(event)
|
||
# pattern 订阅：fnmatch 匹配 event.type（与 InMemoryEventStore 对称）
|
||
for pattern, queues in self._pattern_subscribers.items():
|
||
if fnmatch.fnmatch(event.type, pattern):
|
||
for queue in queues:
|
||
queue.put_nowait(event)
|
||
return event
|
||
|
||
async def _assign_sequence(self, event: Event) -> Event:
|
||
async with self._session_factory() as session:
|
||
result = await session.execute(
|
||
select(func.coalesce(func.max(EventRow.sequence), 0)).where(
|
||
EventRow.run_id == event.run_id
|
||
)
|
||
)
|
||
current = result.scalar_one()
|
||
return event.model_copy(update={"sequence": int(current) + 1})
|
||
|
||
async def list_by_run(self, run_id: str) -> list[Event]:
|
||
async with self._session_factory() as session:
|
||
result = await session.execute(
|
||
select(EventRow).where(EventRow.run_id == run_id).order_by(EventRow.sequence)
|
||
)
|
||
rows = result.scalars().all()
|
||
return [_row_to_event(r) for r in rows]
|
||
|
||
def subscribe(self, run_id: str) -> asyncio.Queue[Event]:
|
||
queue: asyncio.Queue[Event] = asyncio.Queue()
|
||
self._subscribers.setdefault(run_id, []).append(queue)
|
||
return queue
|
||
|
||
def unsubscribe(self, run_id: str, queue: asyncio.Queue[Event]) -> None:
|
||
subs = self._subscribers.get(run_id)
|
||
if subs is not None:
|
||
self._subscribers[run_id] = [q for q in subs if q is not queue]
|
||
|
||
def subscribe_by_pattern(self, pattern: str) -> asyncio.Queue[Event]:
|
||
queue: asyncio.Queue[Event] = asyncio.Queue()
|
||
self._pattern_subscribers.setdefault(pattern, []).append(queue)
|
||
return queue
|
||
|
||
def unsubscribe_by_pattern(self, pattern: str, queue: asyncio.Queue[Event]) -> None:
|
||
subscribers = self._pattern_subscribers.get(pattern)
|
||
if subscribers is not None:
|
||
self._pattern_subscribers[pattern] = [q for q in subscribers if q is not queue]
|
||
|
||
|
||
def _event_to_row(event: Event) -> EventRow:
|
||
"""Event pydantic → events 表行。sequence 未分配时用 0 占位（run 内会被覆盖）。"""
|
||
return EventRow(
|
||
id=event.id,
|
||
type=event.type,
|
||
source=event.source,
|
||
payload_json=json.dumps(event.payload, ensure_ascii=False),
|
||
visibility=event.visibility,
|
||
sensitivity=event.sensitivity,
|
||
run_id=event.run_id,
|
||
character_id=event.character_id,
|
||
trace_id=event.trace_id,
|
||
parent_id=event.parent_id,
|
||
sequence=event.sequence if event.sequence is not None else 0,
|
||
evidence_ids_json=json.dumps(event.evidence_ids, ensure_ascii=False),
|
||
created_at=datetime.fromisoformat(event.created_at),
|
||
)
|
||
|
||
|
||
def _row_to_event(row: EventRow) -> Event:
|
||
"""events 表行 → Event pydantic。"""
|
||
return Event(
|
||
id=row.id,
|
||
type=row.type,
|
||
source=row.source,
|
||
payload=json.loads(row.payload_json),
|
||
visibility=row.visibility,
|
||
sensitivity=row.sensitivity,
|
||
run_id=row.run_id,
|
||
character_id=row.character_id,
|
||
trace_id=row.trace_id,
|
||
parent_id=row.parent_id,
|
||
sequence=row.sequence,
|
||
evidence_ids=json.loads(row.evidence_ids_json),
|
||
created_at=row.created_at.isoformat(),
|
||
)
|