feat：增加分析进度接口

2026-10-10 07:26:24 -07:00
parent e3699d1a30
commit cd5a2b4d8f
15 changed files with 632 additions and 143 deletions
@@ -1,4 +1,4 @@
-"""上传流程：保存原文件 → 子进程解析（超时终止）→ 分配表名 → 保存记录 → 后台用 LLM 分析；以及查看、删除和过期清理"""
+"""上传流程：保存原文件并立即返回 → 后台子进程解析（超时终止，进度可查）→ 分配表名 → 保存记录 → 后台用 LLM 分析；以及查看、删除和过期清理"""

 import asyncio
 import hashlib
@@ -8,7 +8,11 @@ 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
@@ -28,6 +32,7 @@ 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,
@@ -35,11 +40,14 @@ from app.schemas.data_file import (
    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
 # 每小时清理一次目录；一小时内创建的目录不删，避免删掉正在上传的文件
@@ -51,6 +59,8 @@ MAX_NAME_CHARS = 200
 CHUNK_BYTES = 1024 * 1024

 NOT_FOUND = "文件不存在或已过期"
+UPLOAD_NOT_FOUND = "上传任务不存在，可能已中断，请重新上传"
+SAVE_FAILED = "保存文件失败，请稍后重试"
 PARSE_CRASHED = "解析文件时出错，请检查文件内容，或另存为 CSV 后重新上传"


@@ -71,13 +81,11 @@ class DataFileService:
        self._conversations = conversations
        self._ttl = timedelta(seconds=ttl or settings.data_file_ttl)
        self._parsing = asyncio.Semaphore(PARSE_CONCURRENCY)
-        # 文件数检查和写入记录要一起完成，并发上传时才不会超限；单实例部署，进程内的锁就够了
-        self._saving = asyncio.Lock()
+        # 排队和解析中的上传；单实例部署，放进程内存就够了
+        self._uploads = UploadBoard()
        # 上传后的 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 里的时间只精确到毫秒，这里也截到毫秒，刚启动就开始的分析读回后才不会早于它
@@ -86,45 +94,57 @@ class DataFileService:

    async def upload(
        self, user_id: int, conversation_id: str, name: str | None, source: BinaryIO
-    ) -> DataFileRecord:
+    ) -> UploadProgress:
+        """保存原文件并登记解析任务，解析在后台进行，用 get_progress 查进度
+
+        原文件必须在这里写完：请求结束后 source 就被关闭了
+        """
        await self._check_conversation(conversation_id, user_id)
-        # 先粗查一次，满了就不必上传和解析；写入前在锁里再查一次
-        _check_count(await self.store.find(conversation_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)
-            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:
-            # 失败的上传不留文件，会话内的文件数不变
+            # 失败的上传不留文件，也不占名额
+            self._uploads.remove(job)
            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),
+        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,
        )
-        return record
+
+    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
@@ -142,24 +162,6 @@ class DataFileService:
            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,
@@ -233,8 +235,58 @@ class DataFileService:
                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:
-        self._analysing.add(record.file_id)
        task = asyncio.create_task(self._analyse(record))
        self._tasks.add(task)
        task.add_done_callback(self._tasks.discard)
@@ -262,8 +314,6 @@ class DataFileService:
            )
        except Exception as error:
            logger.opt(exception=error).error("保存数据文件分析结果失败")
-        finally:
-            self._analysing.discard(record.file_id)

    def _settle(self, record: DataFileRecord) -> DataFileRecord:
        """服务启动前开始、仍是 pending 的分析已随重启丢失，按失败返回，免得前端一直等
@@ -372,41 +422,100 @@ class DataFileService:
        return removed


-async def _parse(directory: Path) -> ParsedFile:
+ProgressCallback = Callable[[UploadStage, float], None]
+
+
+async def _parse(directory: Path, on_progress: ProgressCallback) -> ParsedFile:
    python, env = _python()
    command = [
        python,
        "-m",
-        "app.datafiles.worker",
+        WORKER_MODULE,
        str(directory / SOURCE_NAME),
        str(directory),
    ]
+    processes: list[subprocess.Popen] = []
    try:
-        # 超时后 subprocess.run 会终止子进程
-        completed = await asyncio.to_thread(
-            subprocess.run,
+        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,
-            capture_output=True,
-            timeout=PARSE_TIMEOUT,
+            stdout=subprocess.PIPE,
+            stderr=stderr,
        )
-    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"])
+        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]:
@@ -426,8 +535,8 @@ 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:
+def _check_count(count: int) -> None:
+    if count >= MAX_FILES_PER_CONVERSATION:
        raise BizException(
            409, f"每个会话最多 {MAX_FILES_PER_CONVERSATION} 个文件，请先删除不用的文件"
        )