570 lines
22 KiB
Python
570 lines
22 KiB
Python
"""上传流程：保存原文件并立即返回 → 后台子进程解析（超时终止，进度可查）→ 分配表名 → 保存记录 → 后台用 LLM 分析；以及查看、删除和过期清理"""
|
||
|
||
import asyncio
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import re
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
from collections.abc import Callable
|
||
from dataclasses import dataclass
|
||
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.datafiles.uploads import UploadBoard, UploadJob
|
||
from app.schemas.data_file import (
|
||
AggregateRequest,
|
||
AggregateResult,
|
||
DataFileRecord,
|
||
DataTable,
|
||
DataTableRows,
|
||
FileInsights,
|
||
UploadProgress,
|
||
UploadStage,
|
||
)
|
||
from app.services.conversation import ConversationStore
|
||
|
||
MAX_FILES_PER_CONVERSATION = 10
|
||
# 解析超时（秒），超时终止子进程；同时最多解析的文件数，其余排队
|
||
WORKER_MODULE = "app.datafiles.worker"
|
||
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 = "文件不存在或已过期"
|
||
UPLOAD_NOT_FOUND = "上传任务不存在，可能已中断，请重新上传"
|
||
SAVE_FAILED = "保存文件失败，请稍后重试"
|
||
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._uploads = UploadBoard()
|
||
# 上传后的 LLM 分析；为空表示不分析，文件的 insights 为空
|
||
self._insights = insights
|
||
self._insight_slots = asyncio.Semaphore(INSIGHTS_CONCURRENCY)
|
||
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
|
||
) -> UploadProgress:
|
||
"""保存原文件并登记解析任务，解析在后台进行，用 get_progress 查进度
|
||
|
||
原文件必须在这里写完：请求结束后 source 就被关闭了
|
||
"""
|
||
await self._check_conversation(conversation_id, user_id)
|
||
records = await self.store.find(conversation_id)
|
||
# 排队和解析中的任务也占名额；检查和登记之间没有 await，并发上传时才不会超限
|
||
_check_count(len(records) + self._uploads.active_count(conversation_id))
|
||
file_id = f"f_{uuid4().hex}"
|
||
job = self._uploads.create(file_id, user_id, conversation_id, _clean_name(name))
|
||
directory = self.root / file_id
|
||
try:
|
||
size, sha256 = await asyncio.to_thread(self._save, source, directory)
|
||
except BaseException:
|
||
# 失败的上传不留文件，也不占名额
|
||
self._uploads.remove(job)
|
||
await asyncio.to_thread(shutil.rmtree, directory, ignore_errors=True)
|
||
raise
|
||
task = asyncio.create_task(self._process(job, directory, size, sha256))
|
||
self._tasks.add(task)
|
||
task.add_done_callback(self._tasks.discard)
|
||
return self._uploads.view(job)
|
||
|
||
async def get_progress(self, user_id: int, file_id: str) -> UploadProgress:
|
||
"""进行中或刚失败的任务看内存里的进度；成功的任务内存里已清掉，按文件记录返回 ready"""
|
||
job = self._uploads.get(file_id)
|
||
if job is not None:
|
||
if job.owner_id != user_id:
|
||
raise BizException(404, UPLOAD_NOT_FOUND)
|
||
return self._uploads.view(job)
|
||
record = await self.store.get(file_id)
|
||
if record is None or record.owner_id != user_id:
|
||
raise BizException(404, UPLOAD_NOT_FOUND)
|
||
return UploadProgress(
|
||
file_id=file_id,
|
||
name=record.name,
|
||
status="ready",
|
||
progress=1.0,
|
||
stage=None,
|
||
eta_seconds=None,
|
||
queued_ahead=0,
|
||
error=None,
|
||
)
|
||
|
||
def list_uploads(self, user_id: int, conversation_id: str) -> list[UploadProgress]:
|
||
"""会话里排队、解析中和刚失败的上传，调用方已校验会话归属"""
|
||
return [
|
||
self._uploads.view(job)
|
||
for job in self._uploads.for_conversation(user_id, conversation_id)
|
||
]
|
||
|
||
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 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, datetime.now(UTC) + self._ttl)
|
||
# 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 sync_screen_files(
|
||
self, user_id: int, screen_id: str, file_ids: set[str]
|
||
) -> None:
|
||
"""大屏保存后调用：它引用的文件固定住，不再引用的放开（重新按保留时长计时）"""
|
||
current = await self.store.pinned_ids(screen_id)
|
||
if added := file_ids - current:
|
||
await self.store.pin(added, user_id, screen_id)
|
||
if removed := current - file_ids:
|
||
await self.store.unpin(screen_id, removed, datetime.now(UTC) + self._ttl)
|
||
|
||
async def release_screen_files(self, screen_id: str) -> None:
|
||
"""大屏进回收站或被彻底删除：它固定的文件全部放开"""
|
||
if current := await self.store.pinned_ids(screen_id):
|
||
await self.store.unpin(screen_id, current, datetime.now(UTC) + self._ttl)
|
||
|
||
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)
|
||
|
||
async def _process(
|
||
self, job: UploadJob, directory: Path, size: int, sha256: str
|
||
) -> None:
|
||
"""后台解析并入库；失败只记在任务上，由进度接口告诉前端"""
|
||
loop = asyncio.get_running_loop()
|
||
|
||
def on_progress(stage: UploadStage, value: float) -> None:
|
||
# 解析在线程里读子进程的输出，转回事件循环再改任务，状态只在一个线程里变
|
||
loop.call_soon_threadsafe(self._uploads.report, job, stage, value)
|
||
|
||
try:
|
||
async with self._parsing:
|
||
self._uploads.start(job)
|
||
parsed = await _parse(directory, on_progress)
|
||
first_number = await self.store.reserve_table_numbers(
|
||
job.conversation_id, len(parsed.tables)
|
||
)
|
||
record = self._record(
|
||
job.file_id,
|
||
job.owner_id,
|
||
job.conversation_id,
|
||
job.name,
|
||
size,
|
||
sha256,
|
||
parsed,
|
||
first_number,
|
||
)
|
||
await self.store.insert(record)
|
||
except BaseException as error:
|
||
# 失败的上传不留文件，会话内的文件数不变
|
||
await asyncio.to_thread(shutil.rmtree, directory, ignore_errors=True)
|
||
if not isinstance(error, Exception):
|
||
# 应用退出时被取消，任务随进程丢失
|
||
self._uploads.remove(job)
|
||
raise
|
||
if isinstance(error, BizException):
|
||
self._uploads.fail(job, error.code, error.msg)
|
||
else:
|
||
logger.opt(exception=error).error("数据文件入库失败：{}", job.file_id)
|
||
self._uploads.fail(job, 500, SAVE_FAILED)
|
||
return
|
||
self._uploads.remove(job)
|
||
if self._insights is not None:
|
||
self._start_insights(record)
|
||
logger.info(
|
||
"数据文件已上传：{} 格式 {}，{} 张表",
|
||
job.file_id,
|
||
record.format,
|
||
len(record.tables),
|
||
)
|
||
|
||
def _start_insights(self, record: DataFileRecord) -> None:
|
||
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("保存数据文件分析结果失败")
|
||
|
||
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
|
||
|
||
|
||
ProgressCallback = Callable[[UploadStage, float], None]
|
||
|
||
|
||
async def _parse(directory: Path, on_progress: ProgressCallback) -> ParsedFile:
|
||
python, env = _python()
|
||
command = [
|
||
python,
|
||
"-m",
|
||
WORKER_MODULE,
|
||
str(directory / SOURCE_NAME),
|
||
str(directory),
|
||
]
|
||
processes: list[subprocess.Popen] = []
|
||
try:
|
||
outcome = await asyncio.to_thread(
|
||
_run_worker, command, env, on_progress, processes
|
||
)
|
||
except asyncio.CancelledError:
|
||
# 线程取消不了，杀掉子进程它才会结束
|
||
for process in processes:
|
||
process.kill()
|
||
raise
|
||
if outcome.timed_out:
|
||
raise BizException(
|
||
422, f"解析超过 {PARSE_TIMEOUT} 秒，文件可能太大或太复杂，请拆分后上传"
|
||
)
|
||
if outcome.returncode != 0 or outcome.result is None:
|
||
logger.error(
|
||
"数据文件解析进程异常退出（{}）：{}",
|
||
outcome.returncode,
|
||
outcome.stderr.decode(errors="replace")[-2000:],
|
||
)
|
||
raise BizException(422, PARSE_CRASHED)
|
||
if error := outcome.result.get("error"):
|
||
raise BizException(error["status"], error["message"])
|
||
return ParsedFile.model_validate(outcome.result["file"])
|
||
|
||
|
||
@dataclass
|
||
class _WorkerOutcome:
|
||
returncode: int
|
||
# 子进程输出的最后一行；没输出结果就退出时为空
|
||
result: dict | None
|
||
stderr: bytes
|
||
timed_out: bool
|
||
|
||
|
||
def _run_worker(
|
||
command: list[str],
|
||
env: dict[str, str] | None,
|
||
on_progress: ProgressCallback,
|
||
processes: list[subprocess.Popen],
|
||
) -> _WorkerOutcome:
|
||
"""运行解析子进程，逐行读它的输出：进度交给 on_progress，最后一行是结果
|
||
|
||
超时由定时器终止子进程，读取随之结束；stderr 写进临时文件，写多了也不会堵住子进程
|
||
"""
|
||
with tempfile.TemporaryFile() as stderr:
|
||
process = subprocess.Popen(
|
||
command,
|
||
cwd=PROJECT_ROOT,
|
||
env=env,
|
||
stdin=subprocess.DEVNULL,
|
||
stdout=subprocess.PIPE,
|
||
stderr=stderr,
|
||
)
|
||
processes.append(process)
|
||
timed_out = threading.Event()
|
||
|
||
def expire() -> None:
|
||
timed_out.set()
|
||
process.kill()
|
||
|
||
timer = threading.Timer(PARSE_TIMEOUT, expire)
|
||
timer.start()
|
||
result = None
|
||
try:
|
||
for line in process.stdout:
|
||
try:
|
||
message = json.loads(line)
|
||
except ValueError:
|
||
continue
|
||
if "progress" in message:
|
||
progress = message["progress"]
|
||
on_progress(progress["stage"], progress["value"])
|
||
else:
|
||
result = message
|
||
returncode = process.wait()
|
||
finally:
|
||
timer.cancel()
|
||
process.kill()
|
||
process.stdout.close()
|
||
stderr.seek(0)
|
||
return _WorkerOutcome(returncode, result, stderr.read(), timed_out.is_set())
|
||
|
||
|
||
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(count: int) -> None:
|
||
if count >= 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
|