174 lines
5.7 KiB
Python
174 lines
5.7 KiB
Python
"""执行模型写的 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], *, threads: int = THREADS
|
||
) -> duckdb.DuckDBPyConnection:
|
||
"""建好视图并锁定的连接；列概况也用它，和模型查询看到的数据一致
|
||
|
||
threads 只给解析子进程算列概况时放宽，模型写的查询一律用 THREADS
|
||
"""
|
||
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 "（无）"
|