61 lines
2.5 KiB
Python
61 lines
2.5 KiB
Python
"""物料目录的 MongoDB 存储：每个组件一条记录，目录版本放在元数据记录中"""
|
|
|
|
from loguru import logger
|
|
from pymongo.asynchronous.database import AsyncDatabase
|
|
|
|
from app.resources.component import load_catalog_file, set_catalog_cache
|
|
from app.schemas.component_resource import ComponentCatalog
|
|
|
|
# 元数据记录的 _id，组件记录的 _id 为组件 key
|
|
META_ID = "__catalog__"
|
|
|
|
|
|
class MongoMaterialRepository:
|
|
def __init__(self, db: AsyncDatabase, prefix: str = ""):
|
|
# 集合名前缀用于测试隔离
|
|
self._materials = db.get_collection(f"{prefix}materials")
|
|
|
|
async def load_all(self) -> ComponentCatalog | None:
|
|
# 没有元数据记录视为尚未导入，由调用方决定回退方式
|
|
meta = await self._materials.find_one({"_id": META_ID})
|
|
if meta is None:
|
|
return None
|
|
documents = await self._materials.find({"_id": {"$ne": META_ID}}).to_list()
|
|
components = [
|
|
{key: value for key, value in item.items() if key not in ("_id", "order")}
|
|
for item in sorted(documents, key=lambda item: item.get("order", 0))
|
|
]
|
|
return ComponentCatalog.model_validate(
|
|
{
|
|
"version": meta["version"],
|
|
"components": components,
|
|
}
|
|
)
|
|
|
|
async def replace_all(self, catalog: ComponentCatalog) -> None:
|
|
# 整体替换目录，供导入脚本使用；order 保持导出文件里的组件顺序
|
|
documents = [
|
|
{
|
|
**component.model_dump(mode="json", by_alias=True),
|
|
"_id": component.key,
|
|
"order": index,
|
|
}
|
|
for index, component in enumerate(catalog.components)
|
|
]
|
|
await self._materials.delete_many({})
|
|
if documents:
|
|
await self._materials.insert_many(documents)
|
|
await self._materials.insert_one({"_id": META_ID, "version": catalog.version})
|
|
|
|
|
|
async def warm_material_cache(db: AsyncDatabase | None) -> ComponentCatalog | None:
|
|
"""启动时把物料目录载入进程缓存；Mongo 未启用或尚未导入时回退读本地 JSON"""
|
|
catalog = await MongoMaterialRepository(db).load_all() if db is not None else None
|
|
if catalog is None:
|
|
catalog = load_catalog_file()
|
|
logger.warning("物料未从 MongoDB 加载，回退使用本地目录文件")
|
|
else:
|
|
logger.info("已从 MongoDB 加载物料 {} 个", len(catalog.components))
|
|
set_catalog_cache(catalog)
|
|
return catalog
|