335 lines
12 KiB
Python
335 lines
12 KiB
Python
"""上传流程：保存原文件 → 子进程解析（超时终止）→ 分配表名 → 保存记录；以及查看、删除和过期清理"""
|
||
|
||
import asyncio
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import re
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
from datetime import UTC, datetime, timedelta
|
||
from pathlib import Path
|
||
from typing import BinaryIO
|
||
from uuid import uuid4
|
||
|
||
from loguru import logger
|
||
|
||
from app.config import PROJECT_ROOT, settings
|
||
from app.core.exceptions import BizException
|
||
from app.datafiles.parse import ParsedFile, ParsedTable
|
||
from app.datafiles.storage import SOURCE_NAME, DataFileStore, file_dir
|
||
from app.schemas.base import JsonObject
|
||
from app.schemas.data_file import (
|
||
DataFileRecord,
|
||
DataTable,
|
||
RequestDataFile,
|
||
RequestDataTable,
|
||
)
|
||
from app.services.conversation import ConversationStore
|
||
|
||
MAX_FILES_PER_CONVERSATION = 10
|
||
# 解析超时（秒），超时终止子进程；同时最多解析的文件数，其余排队
|
||
PARSE_TIMEOUT = 60
|
||
PARSE_CONCURRENCY = 2
|
||
# 每小时清理一次目录；一小时内创建的目录不删，避免删掉正在上传的文件
|
||
SWEEP_INTERVAL = 3600
|
||
SWEEP_GRACE = 3600
|
||
# 只清理这个格式的目录，DATA_FILE_DIR 配错时也不会删到别的东西
|
||
FILE_ID = re.compile(r"^f_[0-9a-f]{32}$")
|
||
MAX_NAME_CHARS = 200
|
||
CHUNK_BYTES = 1024 * 1024
|
||
|
||
NOT_FOUND = "文件不存在或已过期"
|
||
PARSE_CRASHED = "解析文件时出错，请检查文件内容，或另存为 CSV 后重新上传"
|
||
|
||
|
||
class DataFileService:
|
||
def __init__(
|
||
self,
|
||
store: DataFileStore,
|
||
conversations: ConversationStore,
|
||
root: Path | None = None,
|
||
*,
|
||
max_mb: int | None = None,
|
||
ttl: int | None = None,
|
||
):
|
||
self.store = store
|
||
self.root = root or settings.data_file_dir
|
||
self.max_mb = max_mb or settings.data_file_max_mb
|
||
self._conversations = conversations
|
||
self._ttl = timedelta(seconds=ttl or settings.data_file_ttl)
|
||
self._parsing = asyncio.Semaphore(PARSE_CONCURRENCY)
|
||
# 文件数检查、分配表名和写入记录要一起完成，并发上传时才不会超限或重名；单实例部署，进程内的锁就够了
|
||
self._saving = asyncio.Lock()
|
||
|
||
async def upload(
|
||
self, user_id: int, conversation_id: str, name: str | None, source: BinaryIO
|
||
) -> DataFileRecord:
|
||
await self._check_conversation(conversation_id, user_id)
|
||
# 先粗查一次，满了就不必上传和解析；写入前在锁里再查一次
|
||
_check_count(await self.store.find(conversation_id))
|
||
file_id = f"f_{uuid4().hex}"
|
||
directory = file_dir(self.root, file_id)
|
||
try:
|
||
size, sha256 = await asyncio.to_thread(self._save, source, directory)
|
||
async with self._parsing:
|
||
parsed = await _parse(directory)
|
||
async with self._saving:
|
||
existing = await self.store.find(conversation_id, include_deleted=True)
|
||
_check_count([record for record in existing if not record.deleted])
|
||
record = self._record(
|
||
file_id,
|
||
user_id,
|
||
conversation_id,
|
||
_clean_name(name),
|
||
size,
|
||
sha256,
|
||
parsed,
|
||
_next_table_number(existing),
|
||
)
|
||
await self.store.insert(record)
|
||
except BaseException:
|
||
# 失败的上传不留文件，会话内的文件数不变
|
||
await asyncio.to_thread(shutil.rmtree, directory, ignore_errors=True)
|
||
raise
|
||
logger.info(
|
||
"数据文件已上传：{} 格式 {}，{} 张表",
|
||
file_id,
|
||
record.format,
|
||
len(record.tables),
|
||
)
|
||
return record
|
||
|
||
async def list_files(
|
||
self, user_id: int, conversation_id: str
|
||
) -> list[DataFileRecord]:
|
||
await self._check_conversation(conversation_id, user_id)
|
||
records = await self.store.find(conversation_id)
|
||
return [record for record in records if record.owner_id == user_id]
|
||
|
||
async def get_file(self, user_id: int, file_id: str) -> DataFileRecord:
|
||
record = await self.store.get(file_id)
|
||
# 他人的文件与不存在同样返回 404
|
||
if record is None or record.owner_id != user_id:
|
||
raise BizException(404, NOT_FOUND)
|
||
return record
|
||
|
||
async def delete_file(self, user_id: int, file_id: str) -> None:
|
||
await self.get_file(user_id, file_id)
|
||
await self.store.mark_deleted(file_id)
|
||
# Windows 上正在查询的 Parquet 删不掉，留给定时清理
|
||
await asyncio.to_thread(
|
||
shutil.rmtree, file_dir(self.root, file_id), ignore_errors=True
|
||
)
|
||
|
||
async def request_files(
|
||
self, user_id: int, conversation_id: str, attachments: list[JsonObject]
|
||
) -> list[RequestDataFile]:
|
||
"""本次请求可用的文件，调用方已校验会话归属
|
||
|
||
attachments 里 kind 为 data_file 的引用必须是本会话未过期的文件，否则 404；其他种类不处理
|
||
"""
|
||
records = [
|
||
record
|
||
for record in await self.store.find(conversation_id)
|
||
if record.owner_id == user_id
|
||
]
|
||
attached = set()
|
||
for item in attachments:
|
||
if item.get("kind") != "data_file":
|
||
continue
|
||
if not isinstance(file_id := item.get("file_id"), str):
|
||
raise BizException(422, "data_file 附件缺少 file_id")
|
||
attached.add(file_id)
|
||
if attached - {record.file_id for record in records}:
|
||
raise BizException(404, NOT_FOUND)
|
||
return [
|
||
RequestDataFile(
|
||
file_id=record.file_id,
|
||
name=record.name,
|
||
attached_this_turn=record.file_id in attached,
|
||
tables=[
|
||
RequestDataTable.model_validate(
|
||
{
|
||
**table.model_dump(),
|
||
"path": str(
|
||
file_dir(self.root, record.file_id) / table.parquet
|
||
),
|
||
}
|
||
)
|
||
for table in record.tables
|
||
],
|
||
warnings=record.warnings,
|
||
)
|
||
for record in records
|
||
]
|
||
|
||
async def sweep(self) -> int:
|
||
"""删除没有对应记录（过期、已删除、内存模式重启后丢失）的目录，返回删除的个数"""
|
||
live = await self.store.live_ids()
|
||
return await asyncio.to_thread(self._remove_orphans, live)
|
||
|
||
async def run_sweeper(self) -> None:
|
||
while True:
|
||
try:
|
||
if removed := await self.sweep():
|
||
logger.info("已清理 {} 个过期或删除的数据文件目录", removed)
|
||
except Exception as error:
|
||
logger.opt(exception=error).error("清理数据文件目录失败")
|
||
await asyncio.sleep(SWEEP_INTERVAL)
|
||
|
||
async def _check_conversation(self, conversation_id: str, user_id: int) -> None:
|
||
if not await self._conversations.exists(conversation_id, user_id):
|
||
raise BizException(404, "会话不存在或已过期")
|
||
|
||
def _save(self, source: BinaryIO, directory: Path) -> tuple[int, str]:
|
||
"""边写边算大小和 sha256，超过上限立即停止"""
|
||
directory.mkdir(parents=True)
|
||
digest, size = hashlib.sha256(), 0
|
||
with (directory / SOURCE_NAME).open("wb") as target:
|
||
while chunk := source.read(CHUNK_BYTES):
|
||
size += len(chunk)
|
||
if size > self.max_mb * 1024 * 1024:
|
||
raise BizException(413, f"文件不能超过 {self.max_mb} MB")
|
||
digest.update(chunk)
|
||
target.write(chunk)
|
||
return size, digest.hexdigest()
|
||
|
||
def _record(
|
||
self,
|
||
file_id: str,
|
||
user_id: int,
|
||
conversation_id: str,
|
||
name: str,
|
||
size: int,
|
||
sha256: str,
|
||
parsed: ParsedFile,
|
||
first_number: int,
|
||
) -> DataFileRecord:
|
||
tables, warnings = [], list(parsed.warnings)
|
||
for number, table in enumerate(parsed.tables, start=first_number):
|
||
table_name = f"t{number}"
|
||
tables.append(
|
||
DataTable(
|
||
table=table_name,
|
||
label=_label(name, table),
|
||
sheet=table.sheet,
|
||
rows=table.rows,
|
||
columns=table.columns,
|
||
preview=table.preview,
|
||
profile=table.profile,
|
||
parquet=table.path.name,
|
||
)
|
||
)
|
||
warnings.extend(
|
||
warning.model_copy(update={"table": table_name})
|
||
for warning in table.warnings
|
||
)
|
||
now = datetime.now(UTC)
|
||
return DataFileRecord(
|
||
file_id=file_id,
|
||
owner_id=user_id,
|
||
conversation_id=conversation_id,
|
||
name=name,
|
||
format=parsed.format,
|
||
encoding=parsed.encoding,
|
||
size=size,
|
||
sha256=sha256,
|
||
created_at=now,
|
||
expires_at=now + self._ttl,
|
||
tables=tables,
|
||
warnings=warnings,
|
||
)
|
||
|
||
def _remove_orphans(self, live: set[str]) -> int:
|
||
if not self.root.is_dir():
|
||
return 0
|
||
cutoff = time.time() - SWEEP_GRACE
|
||
removed = 0
|
||
for entry in self.root.iterdir():
|
||
if (
|
||
FILE_ID.match(entry.name)
|
||
and entry.name not in live
|
||
and entry.is_dir()
|
||
and entry.stat().st_mtime < cutoff
|
||
):
|
||
shutil.rmtree(entry, ignore_errors=True)
|
||
removed += 1
|
||
return removed
|
||
|
||
|
||
async def _parse(directory: Path) -> ParsedFile:
|
||
python, env = _python()
|
||
command = [
|
||
python,
|
||
"-m",
|
||
"app.datafiles.worker",
|
||
str(directory / SOURCE_NAME),
|
||
str(directory),
|
||
]
|
||
try:
|
||
# 超时后 subprocess.run 会终止子进程
|
||
completed = await asyncio.to_thread(
|
||
subprocess.run,
|
||
command,
|
||
cwd=PROJECT_ROOT,
|
||
env=env,
|
||
stdin=subprocess.DEVNULL,
|
||
capture_output=True,
|
||
timeout=PARSE_TIMEOUT,
|
||
)
|
||
except subprocess.TimeoutExpired:
|
||
raise BizException(
|
||
422, f"解析超过 {PARSE_TIMEOUT} 秒，文件可能太大或太复杂，请拆分后上传"
|
||
) from None
|
||
if completed.returncode != 0:
|
||
logger.error(
|
||
"数据文件解析进程异常退出（{}）：{}",
|
||
completed.returncode,
|
||
completed.stderr.decode(errors="replace")[-2000:],
|
||
)
|
||
raise BizException(422, PARSE_CRASHED)
|
||
result = json.loads(completed.stdout)
|
||
if error := result.get("error"):
|
||
raise BizException(error["status"], error["message"])
|
||
return ParsedFile.model_validate(result["file"])
|
||
|
||
|
||
def _python() -> tuple[str, dict[str, str] | None]:
|
||
"""Windows 的 venv 里 sys.executable 是启动器，由它再拉起真正的解释器；超时时只杀得掉启动器，
|
||
刚启动就被杀时解释器会成为孤儿，一直占着输出管道。和 multiprocessing 一样直接启动基础解释器，
|
||
用 __PYVENV_LAUNCHER__ 告诉它所在的 venv
|
||
"""
|
||
if sys.platform == "win32" and sys._base_executable != sys.executable:
|
||
return sys._base_executable, {
|
||
**os.environ,
|
||
"__PYVENV_LAUNCHER__": sys.executable,
|
||
}
|
||
return sys.executable, None
|
||
|
||
|
||
def _check_count(records: list[DataFileRecord]) -> None:
|
||
if len(records) >= MAX_FILES_PER_CONVERSATION:
|
||
raise BizException(
|
||
409, f"每个会话最多 {MAX_FILES_PER_CONVERSATION} 个文件，请先删除不用的文件"
|
||
)
|
||
|
||
|
||
def _next_table_number(records: list[DataFileRecord]) -> int:
|
||
"""表名在会话内递增，已删除文件的表名也不再使用"""
|
||
used = (int(table.table[1:]) for record in records for table in record.tables)
|
||
return max(used, default=0) + 1
|
||
|
||
|
||
def _clean_name(name: str | None) -> str:
|
||
# 有的浏览器会带上本地路径，只保留文件名
|
||
base = re.split(r"[\\/]", name or "")[-1].strip()
|
||
return base[:MAX_NAME_CHARS] or "未命名文件"
|
||
|
||
|
||
def _label(name: str, table: ParsedTable) -> str:
|
||
label = f"{name} / {table.sheet}" if table.sheet else name
|
||
return f"{label}（{table.part}）" if table.part else label
|