445 lines
17 KiB
Python
445 lines
17 KiB
Python
"""上传流程：保存原文件 → 子进程解析（超时终止）→ 分配表名 → 保存记录 → 后台用 LLM 分析；以及查看、删除和过期清理"""
|
||
|
||
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.aggregate import run_aggregate
|
||
from app.datafiles.insights import (
|
||
INSIGHTS_CONCURRENCY,
|
||
INSIGHTS_TIMEOUT,
|
||
InsightsGenerator,
|
||
)
|
||
from app.datafiles.parse import ParsedFile, ParsedTable
|
||
from app.datafiles.profile import read_rows
|
||
from app.datafiles.query import QueryError
|
||
from app.datafiles.storage import SOURCE_NAME, DataFileStore, table_path
|
||
from app.schemas.data_file import (
|
||
AggregateRequest,
|
||
AggregateResult,
|
||
DataFileRecord,
|
||
DataTable,
|
||
DataTableRows,
|
||
FileInsights,
|
||
)
|
||
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,
|
||
insights: InsightsGenerator | 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()
|
||
# 上传后的 LLM 分析；为空表示不分析，文件的 insights 为空
|
||
self._insights = insights
|
||
self._insight_slots = asyncio.Semaphore(INSIGHTS_CONCURRENCY)
|
||
# 本进程里正在分析的文件，用来拒绝重复的重新分析
|
||
self._analysing: set[str] = set()
|
||
self._tasks: set[asyncio.Task] = set()
|
||
# 在这之前开始、记录里仍是 pending 的分析属于上一个进程，已随重启丢失；
|
||
# MongoDB 里的时间只精确到毫秒，这里也截到毫秒，刚启动就开始的分析读回后才不会早于它
|
||
now = datetime.now(UTC)
|
||
self._started_at = now.replace(microsecond=now.microsecond // 1000 * 1000)
|
||
|
||
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 = 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:
|
||
_check_count(await self.store.find(conversation_id))
|
||
first_number = await self.store.reserve_table_numbers(
|
||
conversation_id, len(parsed.tables)
|
||
)
|
||
record = self._record(
|
||
file_id,
|
||
user_id,
|
||
conversation_id,
|
||
_clean_name(name),
|
||
size,
|
||
sha256,
|
||
parsed,
|
||
first_number,
|
||
)
|
||
await self.store.insert(record)
|
||
except BaseException:
|
||
# 失败的上传不留文件，会话内的文件数不变
|
||
await asyncio.to_thread(shutil.rmtree, directory, ignore_errors=True)
|
||
raise
|
||
if self._insights is not None:
|
||
self._start_insights(record)
|
||
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 [
|
||
self._settle(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 self._settle(record)
|
||
|
||
async def regenerate_insights(self, user_id: int, file_id: str) -> DataFileRecord:
|
||
"""重新分析，用于上次失败之后；正在分析时 409"""
|
||
if self._insights is None:
|
||
raise BizException(422, "未启用文件分析")
|
||
record = await self.get_file(user_id, file_id)
|
||
if file_id in self._analysing:
|
||
raise BizException(409, "正在分析，请稍后")
|
||
self._analysing.add(file_id)
|
||
pending = FileInsights(status="pending", started_at=datetime.now(UTC))
|
||
try:
|
||
await self.store.update_insights(file_id, pending)
|
||
except BaseException:
|
||
self._analysing.discard(file_id)
|
||
raise
|
||
record = record.model_copy(update={"insights": pending})
|
||
self._start_insights(record)
|
||
return record
|
||
|
||
async def get_rows(
|
||
self,
|
||
user_id: int,
|
||
file_id: str,
|
||
table_name: str,
|
||
offset: int,
|
||
limit: int,
|
||
sort_by: str | None = None,
|
||
descending: bool = False,
|
||
) -> DataTableRows:
|
||
"""表的一页数据；表或排序列不存在时 404 / 422，文件已被清理时同样 404"""
|
||
record = await self.get_file(user_id, file_id)
|
||
table, path = self._table(record, table_name)
|
||
if sort_by is not None and sort_by not in {c.name for c in table.columns}:
|
||
raise BizException(422, f"排序列不存在：{sort_by}")
|
||
rows = await asyncio.to_thread(
|
||
read_rows, path, offset, limit, sort_by, descending
|
||
)
|
||
return DataTableRows(total=table.rows, offset=offset, rows=rows)
|
||
|
||
async def aggregate(
|
||
self, user_id: int, file_id: str, table_name: str, request: AggregateRequest
|
||
) -> AggregateResult:
|
||
"""对一张表分组汇总，供前端做图表数据；和 agent 的 query_data 同一套只读沙箱"""
|
||
record = await self.get_file(user_id, file_id)
|
||
table, path = self._table(record, table_name)
|
||
try:
|
||
return await asyncio.to_thread(run_aggregate, path, request, table.columns)
|
||
except QueryError as error:
|
||
raise BizException(422, str(error)) from None
|
||
|
||
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, self.root / file_id, ignore_errors=True)
|
||
|
||
async def request_files(
|
||
self, user_id: int, conversation_id: str, attached: set[str]
|
||
) -> list[DataFileRecord]:
|
||
"""本次请求可用的文件：当前用户在本会话里未过期的文件，调用方已校验会话归属
|
||
|
||
attached 是本轮附件引用的文件，必须都在其中，否则 404
|
||
"""
|
||
records = [
|
||
record
|
||
for record in await self.store.find(conversation_id)
|
||
if record.owner_id == user_id
|
||
]
|
||
if attached - {record.file_id for record in records}:
|
||
raise BizException(404, NOT_FOUND)
|
||
return records
|
||
|
||
async def aclose(self) -> None:
|
||
"""应用退出时取消还在进行的分析，记录里留下的 pending 下次启动后读取时按失败返回"""
|
||
for task in list(self._tasks):
|
||
task.cancel()
|
||
await asyncio.gather(*self._tasks, return_exceptions=True)
|
||
|
||
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)
|
||
|
||
def _start_insights(self, record: DataFileRecord) -> None:
|
||
self._analysing.add(record.file_id)
|
||
task = asyncio.create_task(self._analyse(record))
|
||
self._tasks.add(task)
|
||
task.add_done_callback(self._tasks.discard)
|
||
|
||
async def _analyse(self, record: DataFileRecord) -> None:
|
||
"""后台分析并写回记录；任何失败都只记为 failed，不影响文件使用"""
|
||
assert self._insights is not None
|
||
try:
|
||
try:
|
||
async with self._insight_slots:
|
||
insights = await asyncio.wait_for(
|
||
self._insights(record, self.root), INSIGHTS_TIMEOUT
|
||
)
|
||
except Exception as error:
|
||
logger.opt(exception=error).warning(
|
||
"数据文件分析失败：{}", record.file_id
|
||
)
|
||
insights = FileInsights(status="failed")
|
||
await self.store.update_insights(record.file_id, insights)
|
||
logger.info(
|
||
"数据文件分析完成：{} {}，{} 条推荐",
|
||
record.file_id,
|
||
insights.status,
|
||
len(insights.suggestions),
|
||
)
|
||
except Exception as error:
|
||
logger.opt(exception=error).error("保存数据文件分析结果失败")
|
||
finally:
|
||
self._analysing.discard(record.file_id)
|
||
|
||
def _settle(self, record: DataFileRecord) -> DataFileRecord:
|
||
"""服务启动前开始、仍是 pending 的分析已随重启丢失，按失败返回，免得前端一直等
|
||
|
||
只看记录本身，不写回存储：读到的记录可能已经过时，写回会覆盖刚完成的结果
|
||
"""
|
||
insights = record.insights
|
||
if (
|
||
insights is not None
|
||
and insights.status == "pending"
|
||
and (insights.started_at is None or insights.started_at < self._started_at)
|
||
):
|
||
return record.model_copy(update={"insights": FileInsights(status="failed")})
|
||
return record
|
||
|
||
def _table(self, record: DataFileRecord, name: str) -> tuple[DataTable, Path]:
|
||
"""文件里的一张表和它的 Parquet
|
||
|
||
表不存在，或 Parquet 已被删除（删除时目录没删净、过期清理）时 404
|
||
"""
|
||
table = next((item for item in record.tables if item.table == name), None)
|
||
if table is None:
|
||
raise BizException(404, "数据表不存在")
|
||
path = table_path(self.root, record.file_id, table)
|
||
if not path.is_file():
|
||
raise BizException(404, NOT_FOUND)
|
||
return table, path
|
||
|
||
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 file_too_large(self.max_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),
|
||
**table.model_dump(exclude={"part", "warnings"}),
|
||
)
|
||
)
|
||
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,
|
||
insights=(
|
||
None
|
||
if self._insights is None
|
||
else FileInsights(status="pending", started_at=now)
|
||
),
|
||
)
|
||
|
||
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 file_too_large(max_mb: int) -> BizException:
|
||
return BizException(413, f"文件不能超过 {max_mb} MB")
|
||
|
||
|
||
def _check_count(records: list[DataFileRecord]) -> None:
|
||
if len(records) >= MAX_FILES_PER_CONVERSATION:
|
||
raise BizException(
|
||
409, f"每个会话最多 {MAX_FILES_PER_CONVERSATION} 个文件，请先删除不用的文件"
|
||
)
|
||
|
||
|
||
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
|