feat：完成大部分数据文件上传解析功能
This commit is contained in:
@@ -0,0 +1,334 @@
|
||||
"""上传流程：保存原文件 → 子进程解析（超时终止）→ 分配表名 → 保存记录；以及查看、删除和过期清理"""
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user