239 lines
8.6 KiB
Python
239 lines
8.6 KiB
Python
"""大屏存储的 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})
|