297 lines
11 KiB
Python
297 lines
11 KiB
Python
"""数据文件的存储：文件放本地目录，元数据启用 MongoDB 时进集合 data_files，否则放进程内存
|
|
|
|
目录结构是 DATA_FILE_DIR/<file_id>/，里面有原文件 source 和每张表的 1.parquet、2.parquet……
|
|
DuckDB 只能读本地文件，所以文件不进 MongoDB；单实例部署，本地磁盘就够了
|
|
会话内的表名 t1、t2……由每个会话一个的计数器分配，文件删除、过期后表名也不再使用
|
|
"""
|
|
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any, Protocol
|
|
|
|
from bson.codec_options import CodecOptions
|
|
from pymongo import ASCENDING, ReturnDocument
|
|
|
|
from app.schemas.data_file import DataFileRecord, DataTable, FileInsights
|
|
|
|
# 原文件不带扩展名：calamine 见到 .xls 等扩展名会按扩展名打开，用户的扩展名和内容不符时会失败
|
|
SOURCE_NAME = "source"
|
|
|
|
|
|
# 被大屏固定的文件的占位过期时间：MongoDB 的 TTL 索引只认 expires_at，推到遥远的将来就不会被删
|
|
PINNED_EXPIRES = datetime(9999, 1, 1, tzinfo=UTC)
|
|
|
|
|
|
def table_path(root: Path, file_id: str, table: DataTable) -> Path:
|
|
return root / file_id / table.parquet
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(UTC)
|
|
|
|
|
|
class DataFileStore(Protocol):
|
|
"""内存和 MongoDB 实现同一组方法；过期的记录一律当作不存在"""
|
|
|
|
async def insert(self, record: DataFileRecord) -> None: ...
|
|
|
|
async def get(self, file_id: str) -> DataFileRecord | None:
|
|
"""未过期且未删除的记录"""
|
|
...
|
|
|
|
async def find(self, conversation_id: str) -> list[DataFileRecord]:
|
|
"""会话内未过期且未删除的记录，按上传时间排列"""
|
|
...
|
|
|
|
async def reserve_table_numbers(self, conversation_id: str, count: int) -> int:
|
|
"""为会话预留 count 个连续的表号，返回第一个；表号从 1 开始，不会重复分配"""
|
|
...
|
|
|
|
async def mark_deleted(
|
|
self, file_id: str, expires_at: datetime | None = None
|
|
) -> None:
|
|
"""给了 expires_at 就同时撤销所有大屏对它的固定，并把过期时间改成它（被固定的文件不能停在 9999 年）"""
|
|
...
|
|
|
|
async def update_insights(self, file_id: str, insights: FileInsights) -> None:
|
|
"""记录已过期或不存在时什么都不做"""
|
|
...
|
|
|
|
async def live_ids(self) -> set[str]:
|
|
"""未过期且未删除的 file_id，清理目录时用"""
|
|
...
|
|
|
|
async def pin(self, file_ids: set[str], owner_id: int, screen_id: str) -> None:
|
|
"""让 screen_id 固定这些文件（只对 owner 名下、未过期、未删除的生效），被固定的文件不再过期"""
|
|
...
|
|
|
|
async def unpin(
|
|
self, screen_id: str, file_ids: set[str], expires_at: datetime
|
|
) -> None:
|
|
"""撤销 screen_id 对这些文件的固定；不再被任何大屏固定的文件，过期时间改为 expires_at"""
|
|
...
|
|
|
|
async def pinned_ids(self, screen_id: str) -> set[str]:
|
|
"""被 screen_id 固定的文件，包括已删除的（它们也要撤销固定）"""
|
|
...
|
|
|
|
|
|
class InMemoryDataFileStore:
|
|
def __init__(self) -> None:
|
|
self._records: dict[str, DataFileRecord] = {}
|
|
# 会话 → 已分配的最大表号
|
|
self._table_numbers: dict[str, int] = {}
|
|
|
|
def _purge_expired(self) -> None:
|
|
now = _now()
|
|
for file_id in [
|
|
key for key, record in self._records.items() if record.expires_at <= now
|
|
]:
|
|
del self._records[file_id]
|
|
|
|
async def insert(self, record: DataFileRecord) -> None:
|
|
self._purge_expired()
|
|
self._records[record.file_id] = record
|
|
|
|
async def get(self, file_id: str) -> DataFileRecord | None:
|
|
self._purge_expired()
|
|
record = self._records.get(file_id)
|
|
return None if record is None or record.deleted else record
|
|
|
|
async def find(self, conversation_id: str) -> list[DataFileRecord]:
|
|
self._purge_expired()
|
|
records = [
|
|
record
|
|
for record in self._records.values()
|
|
if record.conversation_id == conversation_id and not record.deleted
|
|
]
|
|
return sorted(records, key=lambda record: record.created_at)
|
|
|
|
async def reserve_table_numbers(self, conversation_id: str, count: int) -> int:
|
|
first = self._table_numbers.get(conversation_id, 0) + 1
|
|
self._table_numbers[conversation_id] = first + count - 1
|
|
return first
|
|
|
|
async def mark_deleted(
|
|
self, file_id: str, expires_at: datetime | None = None
|
|
) -> None:
|
|
if (record := self._records.get(file_id)) is not None:
|
|
changes: dict = {"deleted": True}
|
|
if expires_at is not None:
|
|
changes.update(pinned_by=[], expires_at=expires_at)
|
|
self._records[file_id] = record.model_copy(update=changes)
|
|
|
|
async def update_insights(self, file_id: str, insights: FileInsights) -> None:
|
|
if (record := self._records.get(file_id)) is not None:
|
|
self._records[file_id] = record.model_copy(update={"insights": insights})
|
|
|
|
async def live_ids(self) -> set[str]:
|
|
self._purge_expired()
|
|
return {key for key, record in self._records.items() if not record.deleted}
|
|
|
|
async def pin(self, file_ids, owner_id, screen_id):
|
|
self._purge_expired()
|
|
for file_id in file_ids:
|
|
record = self._records.get(file_id)
|
|
if record is None or record.deleted or record.owner_id != owner_id:
|
|
continue
|
|
pinned_by = record.pinned_by
|
|
if screen_id not in pinned_by:
|
|
pinned_by = [*pinned_by, screen_id]
|
|
self._records[file_id] = record.model_copy(
|
|
update={"pinned_by": pinned_by, "expires_at": PINNED_EXPIRES}
|
|
)
|
|
|
|
async def unpin(self, screen_id, file_ids, expires_at):
|
|
for file_id in file_ids:
|
|
record = self._records.get(file_id)
|
|
if record is None or screen_id not in record.pinned_by:
|
|
continue
|
|
pinned_by = [item for item in record.pinned_by if item != screen_id]
|
|
changes = {"pinned_by": pinned_by}
|
|
if not pinned_by:
|
|
changes["expires_at"] = expires_at
|
|
self._records[file_id] = record.model_copy(update=changes)
|
|
|
|
async def pinned_ids(self, screen_id):
|
|
return {
|
|
key
|
|
for key, record in self._records.items()
|
|
if screen_id in record.pinned_by
|
|
}
|
|
|
|
|
|
class MongoDataFileStore:
|
|
def __init__(self, db: Any, prefix: str = ""):
|
|
# 读回的时间统一带 UTC 时区；集合名前缀用于测试隔离
|
|
options = CodecOptions(tz_aware=True, tzinfo=UTC)
|
|
self._files = db.get_collection(f"{prefix}data_files", codec_options=options)
|
|
# 每个会话一条：_id 是会话 ID，tables 是已分配的最大表号；会话一直保留，计数器也不过期
|
|
self._table_numbers = db.get_collection(f"{prefix}data_file_table_numbers")
|
|
|
|
async def ensure_indexes(self) -> None:
|
|
# 过期记录由 TTL 索引删除；TTL 每分钟才扫一次，查询时仍按 expires_at 过滤
|
|
await self._files.create_index("expires_at", expireAfterSeconds=0)
|
|
await self._files.create_index(
|
|
[("conversation_id", ASCENDING), ("created_at", ASCENDING)]
|
|
)
|
|
# 按大屏找它固定的文件
|
|
await self._files.create_index("pinned_by")
|
|
|
|
async def insert(self, record: DataFileRecord) -> None:
|
|
await self._files.insert_one({"_id": record.file_id, **record.model_dump()})
|
|
|
|
async def get(self, file_id: str) -> DataFileRecord | None:
|
|
document = await self._files.find_one(
|
|
{"_id": file_id, "deleted": False, "expires_at": {"$gt": _now()}}
|
|
)
|
|
return None if document is None else _record(document)
|
|
|
|
async def find(self, conversation_id: str) -> list[DataFileRecord]:
|
|
query = {
|
|
"conversation_id": conversation_id,
|
|
"deleted": False,
|
|
"expires_at": {"$gt": _now()},
|
|
}
|
|
documents = await (
|
|
self._files.find(query)
|
|
.sort([("created_at", ASCENDING), ("_id", ASCENDING)])
|
|
.to_list()
|
|
)
|
|
return [_record(document) for document in documents]
|
|
|
|
async def reserve_table_numbers(self, conversation_id: str, count: int) -> int:
|
|
document = await self._table_numbers.find_one_and_update(
|
|
{"_id": conversation_id},
|
|
{"$inc": {"tables": count}},
|
|
upsert=True,
|
|
return_document=ReturnDocument.AFTER,
|
|
)
|
|
return document["tables"] - count + 1
|
|
|
|
async def mark_deleted(
|
|
self, file_id: str, expires_at: datetime | None = None
|
|
) -> None:
|
|
changes: dict = {"deleted": True}
|
|
if expires_at is not None:
|
|
changes.update(pinned_by=[], expires_at=expires_at)
|
|
await self._files.update_one({"_id": file_id}, {"$set": changes})
|
|
|
|
async def update_insights(self, file_id: str, insights: FileInsights) -> None:
|
|
await self._files.update_one(
|
|
{"_id": file_id}, {"$set": {"insights": insights.model_dump()}}
|
|
)
|
|
|
|
async def live_ids(self) -> set[str]:
|
|
return set(
|
|
await self._files.distinct(
|
|
"_id", {"deleted": False, "expires_at": {"$gt": _now()}}
|
|
)
|
|
)
|
|
|
|
async def pin(self, file_ids, owner_id, screen_id):
|
|
if not file_ids:
|
|
return
|
|
await self._files.update_many(
|
|
{
|
|
"_id": {"$in": list(file_ids)},
|
|
"owner_id": owner_id,
|
|
"deleted": False,
|
|
"expires_at": {"$gt": _now()},
|
|
},
|
|
{
|
|
"$addToSet": {"pinned_by": screen_id},
|
|
"$set": {"expires_at": PINNED_EXPIRES},
|
|
},
|
|
)
|
|
|
|
async def unpin(self, screen_id, file_ids, expires_at):
|
|
if not file_ids:
|
|
return
|
|
# 一条聚合管道更新：去掉这个大屏，且只有数组因此变空时才重新计时。
|
|
# 单文档原子，不会停在"数组空了但仍是 9999 年"；只改确实固定了它的文档
|
|
await self._files.update_many(
|
|
{"_id": {"$in": list(file_ids)}, "pinned_by": screen_id},
|
|
[
|
|
{
|
|
"$set": {
|
|
"pinned_by": {
|
|
"$filter": {
|
|
"input": "$pinned_by",
|
|
"cond": {"$ne": ["$$this", screen_id]},
|
|
}
|
|
}
|
|
}
|
|
},
|
|
{
|
|
"$set": {
|
|
"expires_at": {
|
|
"$cond": [
|
|
{"$eq": [{"$size": "$pinned_by"}, 0]},
|
|
expires_at,
|
|
"$expires_at",
|
|
]
|
|
}
|
|
}
|
|
},
|
|
],
|
|
)
|
|
|
|
async def pinned_ids(self, screen_id):
|
|
return set(await self._files.distinct("_id", {"pinned_by": screen_id}))
|
|
|
|
|
|
def _record(document: dict[str, Any]) -> DataFileRecord:
|
|
return DataFileRecord.model_validate(
|
|
{key: value for key, value in document.items() if key != "_id"}
|
|
)
|
|
|
|
|
|
async def open_data_file_store(db: Any) -> DataFileStore:
|
|
"""db 为空（未启用 MongoDB）时用内存实现，和会话记录一致"""
|
|
if db is None:
|
|
return InMemoryDataFileStore()
|
|
store = MongoDataFileStore(db)
|
|
await store.ensure_indexes()
|
|
return store
|