feat：完成大部分数据文件上传解析功能

2026-10-08 08:23:01 -07:00
parent 05c050d2ba
commit bb6a9405db
27 changed files with 2675 additions and 18 deletions
@@ -0,0 +1,168 @@
+"""执行模型写的 SQL：连接只能读本次可见的 Parquet 文件，只允许单条 SELECT
+
+每次查询新建一个内存连接：
+- allowed_paths 只放这些 Parquet 文件；不用 allowed_directories，否则同目录下的其他文件能被 glob 读到，
+  COPY、ATTACH 也能往目录里写
+- 关闭外部访问、扩展加载和 Python 变量替换扫描，最后锁定配置，SQL 里的 SET 一律失败
+"""
+
+import math
+import threading
+from collections.abc import Mapping
+from dataclasses import dataclass
+from datetime import date, datetime, time, timedelta
+from decimal import Decimal
+from pathlib import Path
+from typing import Any
+
+import duckdb
+
+MEMORY_LIMIT = "512MB"
+THREADS = 2
+MAX_ROWS = 200
+MAX_CELL_CHARS = 200
+
+
+class QueryError(ValueError):
+    """查询失败的原因，原样返回给模型修正"""
+
+
+@dataclass(frozen=True)
+class QueryResult:
+    columns: list[str]
+    rows: list[list[Any]]
+    # 结果超过 max_rows 行，只返回了前 max_rows 行
+    truncated: bool
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "columns": self.columns,
+            "rows": self.rows,
+            "row_count": len(self.rows),
+            "truncated": self.truncated,
+        }
+
+
+def run_query(
+    sql: str,
+    tables: Mapping[str, Path],
+    *,
+    timeout: float,
+    max_rows: int = MAX_ROWS,
+) -> QueryResult:
+    """tables 是表名到 Parquet 文件的映射，SQL 里只能用这些表名"""
+    try:
+        connection = connect(tables)
+    except duckdb.Error as error:
+        # 任务执行期间文件被删除或过期清理
+        raise QueryError(
+            f"数据表无法读取，文件可能已被删除：{_first_line(error)}"
+        ) from None
+    try:
+        _check_statement(connection, sql, tables)
+        timer = threading.Timer(timeout, connection.interrupt)
+        timer.start()
+        try:
+            connection.execute(sql)
+            columns = [column[0] for column in connection.description]
+            rows = connection.fetchmany(max_rows + 1)
+        except duckdb.InterruptException:
+            raise QueryError(
+                f"查询超过 {timeout:g} 秒被中止，请缩小范围或先聚合"
+            ) from None
+        except duckdb.PermissionException:
+            raise QueryError(f"只能查询上传的数据表：{_table_list(tables)}") from None
+        except duckdb.Error as error:
+            raise QueryError(
+                f"{_first_line(error)}；可用的表：{_table_list(tables)}"
+            ) from None
+        finally:
+            timer.cancel()
+    finally:
+        connection.close()
+    return QueryResult(
+        columns,
+        [[json_value(value) for value in row] for row in rows[:max_rows]],
+        len(rows) > max_rows,
+    )
+
+
+def connect(tables: Mapping[str, Path]) -> duckdb.DuckDBPyConnection:
+    """建好视图并锁定的连接；列概况也用它，和模型查询看到的数据一致"""
+    connection = duckdb.connect(
+        ":memory:",
+        config={
+            "autoinstall_known_extensions": False,
+            "autoload_known_extensions": False,
+            "python_enable_replacements": False,
+            "memory_limit": MEMORY_LIMIT,
+            "threads": THREADS,
+        },
+    )
+    try:
+        paths = [str(Path(path).resolve()) for path in tables.values()]
+        for name, path in zip(tables, paths):
+            connection.execute(
+                f"CREATE VIEW {quote_identifier(name)} AS FROM read_parquet({_literal(path)})"
+            )
+        if paths:
+            connection.execute("SET allowed_paths = ?", [paths])
+        connection.execute("SET enable_external_access = false")
+        connection.execute("SET lock_configuration = true")
+    except BaseException:
+        connection.close()
+        raise
+    return connection
+
+
+def quote_identifier(name: str) -> str:
+    return '"' + name.replace('"', '""') + '"'
+
+
+def json_value(value: Any, max_chars: int = MAX_CELL_CHARS) -> Any:
+    """查询结果转成 JSON 能表示的值；长文本截断"""
+    if value is None or isinstance(value, (bool, int)):
+        return value
+    if isinstance(value, float):
+        return None if math.isnan(value) or math.isinf(value) else value
+    if isinstance(value, Decimal):
+        return float(value)
+    if isinstance(value, str):
+        return value if len(value) <= max_chars else value[:max_chars] + "…"
+    if isinstance(value, (datetime, date, time)):
+        return value.isoformat()
+    if isinstance(value, (list, tuple)):
+        return [json_value(item, max_chars) for item in value]
+    if isinstance(value, dict):
+        return {str(key): json_value(item, max_chars) for key, item in value.items()}
+    if isinstance(value, timedelta):
+        return str(value)
+    return json_value(str(value), max_chars)
+
+
+def _check_statement(
+    connection: duckdb.DuckDBPyConnection, sql: str, tables: Mapping[str, Path]
+) -> None:
+    try:
+        statements = connection.extract_statements(sql)
+    except duckdb.Error as error:
+        raise QueryError(
+            f"{_first_line(error)}；可用的表：{_table_list(tables)}"
+        ) from None
+    # PIVOT 不写 IN 时会先建一个枚举类型，解析出来是两条语句
+    if len(statements) != 1 or statements[0].type != duckdb.StatementType.SELECT:
+        raise QueryError(
+            "只能执行一条 SELECT 查询（可以用 WITH、UNPIVOT；PIVOT 要用 IN 写明取值）"
+        )
+
+
+def _literal(text: str) -> str:
+    return "'" + text.replace("'", "''") + "'"
+
+
+def _first_line(error: Exception) -> str:
+    return str(error).strip().splitlines()[0]
+
+
+def _table_list(tables: Mapping[str, Path]) -> str:
+    return "、".join(tables) if tables else "（无）"