"""大屏存储的 MongoDB 实现，方法语义与内存实现一致 集合：screens（草稿，一屏一条）、screen_publishes（发布快照，只增不改） """ import re from contextlib import contextmanager from datetime import UTC from typing import Any from bson.codec_options import CodecOptions from bson.errors import InvalidDocument from pydantic import BaseModel from pymongo import ASCENDING, DESCENDING, ReturnDocument from pymongo.errors import DocumentTooLarge, DuplicateKeyError from app.core.exceptions import BizException from app.schemas.saved_screen import ( PublishEntry, PublishSnapshot, ScreenRecord, ScreenSummary, ) from app.services.screen_store import canvas_size # 列表和元数据更新不需要页面，只取画布；Mongo 的投影不能同时"排除 page"和"包含 page.canvas"，所以列字段 SUMMARY_PROJECTION = { "owner_id": 1, "name": 1, "thumbnail": 1, "page.canvas": 1, "revision": 1, "created_at": 1, "updated_at": 1, "deleted_at": 1, "publish": 1, } @contextmanager def _storable(): """页面里的超大整数、含 NUL 的字段名、编码后超过 16 MB 的文档，驱动会抛各种非业务异常 10 MB 的 JSON 编码成 BSON 后会涨（数组的下标要当字符串键存），所以体积上限靠这里兜底 """ try: yield except (OverflowError, InvalidDocument, DocumentTooLarge): raise BizException( 422, "大屏数据无法保存：数字超出范围、字段名不合法或体积过大" ) from None # update 自己拼的条件，调用方的 where 不许覆盖 RESERVED_FILTER_KEYS = {"_id", "owner_id", "deleted_at", "revision"} def _owned(screen_id: str, owner_id: int, deleted: bool) -> dict[str, Any]: return { "_id": screen_id, "owner_id": owner_id, "deleted_at": {"$ne": None} if deleted else None, } def _encode(value: Any) -> Any: return value.model_dump() if isinstance(value, BaseModel) else value def _summary(document: dict[str, Any]) -> ScreenSummary: canvas = (document.get("page") or {}).get("canvas") data = {k: v for k, v in document.items() if k not in ("_id", "page")} return ScreenSummary.model_validate( {**data, "screen_id": document["_id"], "canvas": canvas_size(canvas)} ) class MongoScreenStore: def __init__(self, db: Any, prefix: str = ""): # 读回的时间统一带 UTC 时区；集合名前缀用于测试隔离 options = CodecOptions(tz_aware=True, tzinfo=UTC) self._screens = db.get_collection(f"{prefix}screens", codec_options=options) self._publishes = db.get_collection( f"{prefix}screen_publishes", codec_options=options ) async def ensure_indexes(self) -> None: await self._screens.create_index( [ ("owner_id", ASCENDING), ("deleted_at", ASCENDING), ("updated_at", DESCENDING), ] ) # 分享链接唯一；没发布过的大屏没有 publish.share_id，不进索引 await self._screens.create_index( "publish.share_id", unique=True, partialFilterExpression={"publish.share_id": {"$type": "string"}}, ) await self._publishes.create_index( [("screen_id", ASCENDING), ("version", DESCENDING)], unique=True ) # 回收站过期清理按 deleted_at 跨用户扫描；没删除的不进索引 await self._screens.create_index( "deleted_at", partialFilterExpression={"deleted_at": {"$type": "date"}} ) async def insert(self, record: ScreenRecord) -> None: document = record.model_dump(exclude={"screen_id", "canvas"}) with _storable(): await self._screens.insert_one({"_id": record.screen_id, **document}) async def get(self, screen_id, owner_id, *, deleted=False): document = await self._screens.find_one(_owned(screen_id, owner_id, deleted)) if document is None: return None return ScreenRecord.model_validate( {**_summary(document).model_dump(), "page": document["page"]} ) async def get_summary(self, screen_id, owner_id, *, deleted=False): document = await self._screens.find_one( _owned(screen_id, owner_id, deleted), SUMMARY_PROJECTION ) return None if document is None else _summary(document) async def update( self, screen_id, owner_id, changes, *, revision=None, where=None, deleted=False ): query = _owned(screen_id, owner_id, deleted) if revision is not None: query["revision"] = revision if clash := RESERVED_FILTER_KEYS & (where or {}).keys(): raise ValueError(f"where 不能覆盖所有权和版本条件：{sorted(clash)}") query.update(where or {}) with _storable(): document = await self._screens.find_one_and_update( query, {"$set": {key: _encode(value) for key, value in changes.items()}}, projection=SUMMARY_PROJECTION, return_document=ReturnDocument.AFTER, ) return None if document is None else _summary(document) async def search( self, owner_id, *, keyword, status, sort, descending, limit, offset, deleted ): query: dict[str, Any] = { "owner_id": owner_id, "deleted_at": {"$ne": None} if deleted else None, } if keyword: query["name"] = {"$regex": re.escape(keyword), "$options": "i"} if status == "draft": query["publish"] = None elif status is not None: query["publish.status"] = status direction = DESCENDING if descending else ASCENDING documents = await ( self._screens.find(query, SUMMARY_PROJECTION) .sort([(sort, direction), ("_id", direction)]) .skip(offset) .limit(limit) .to_list() ) return [_summary(d) for d in documents], await self._screens.count_documents( query ) async def get_by_share_id(self, share_id): document = await self._screens.find_one( {"publish.share_id": share_id, "deleted_at": None}, SUMMARY_PROJECTION ) return None if document is None else _summary(document) async def list_expired_trash(self, deleted_before, limit): documents = await ( self._screens.find( # 带 $type 才蕴含部分索引的过滤条件，查询才用得上索引 {"deleted_at": {"$type": "date", "$lt": deleted_before}}, SUMMARY_PROJECTION, ) .sort("deleted_at", ASCENDING) .limit(limit) .to_list() ) return [_summary(d) for d in documents] async def purge(self, screen_id, owner_id): # 带"仍在回收站"的条件：清理期间被还原的大屏不能删 result = await self._screens.delete_one(_owned(screen_id, owner_id, True)) return result.deleted_count == 1 async def delete_publishes(self, screen_id): await self._publishes.delete_many({"screen_id": screen_id}) async def last_publish_version(self, screen_id): document = await self._publishes.find_one( {"screen_id": screen_id}, {"version": 1}, sort=[("version", DESCENDING)] ) return 0 if document is None else document["version"] async def insert_publish(self, snapshot): try: with _storable(): await self._publishes.insert_one(snapshot.model_dump()) except DuplicateKeyError: return False return True async def get_publish(self, screen_id, version): document = await self._publishes.find_one( {"screen_id": screen_id, "version": version} ) if document is None: return None return PublishSnapshot.model_validate( {k: v for k, v in document.items() if k != "_id"} ) async def list_publishes(self, screen_id, limit, offset): query = {"screen_id": screen_id} documents = await ( self._publishes.find(query, {"page": 0}) .sort("version", DESCENDING) .skip(offset) .limit(limit) .to_list() ) entries = [ PublishEntry.model_validate({k: v for k, v in d.items() if k != "_id"}) for d in documents ] return entries, await self._publishes.count_documents(query) async def delete_publish(self, screen_id, version): await self._publishes.delete_one({"screen_id": screen_id, "version": version})