107 lines
4.7 KiB
Python
107 lines
4.7 KiB
Python
"""DuckDB: Dateien umwandeln und schreibgeschützt abfragen (→ Konzept 7.7).
|
|
|
|
Abfragen laufen in einer In-Memory-Datenbank, an die die Upload-Dateien des
|
|
Workspace schreibgeschützt angehängt sind. Danach wird jeder Zugriff auf Dateien
|
|
und Netz gesperrt und die Konfiguration festgeschrieben; Erweiterungen werden
|
|
weder installiert noch geladen. Zugelassen ist genau eine SELECT-Anweisung.
|
|
"""
|
|
|
|
import threading
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import duckdb
|
|
|
|
from app.core.config import Settings
|
|
from app.tools.registry import ToolError
|
|
|
|
READERS = {"csv": "read_csv($1, auto_detect=true, sample_size=20000)",
|
|
"json": "read_json_auto($1)",
|
|
"parquet": "read_parquet($1)"}
|
|
BASE_CONFIG = {"autoinstall_known_extensions": False, "autoload_known_extensions": False, "threads": 2}
|
|
CELL_CHARS = 200
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Table:
|
|
name: str # Name für SQL
|
|
path: Path # DuckDB-Datei mit Tabelle data
|
|
|
|
|
|
def convert(source: Path, fmt: str, target: Path, settings: Settings) -> tuple[list[dict[str, str]], int]:
|
|
"""Datei in eine DuckDB-Datei mit einer Tabelle data umwandeln; liefert Spalten und Zeilenzahl."""
|
|
target.unlink(missing_ok=True)
|
|
try:
|
|
with duckdb.connect(str(target), config={**BASE_CONFIG,
|
|
"memory_limit": f"{settings.query_memory_mb}MB"}) as con:
|
|
con.execute(f"CREATE TABLE data AS SELECT * FROM {READERS[fmt]}", [str(source)])
|
|
columns = [{"name": name, "type": kind} for name, kind, *_ in con.execute("DESCRIBE data").fetchall()]
|
|
rows = con.execute("SELECT count(*) FROM data").fetchone()[0]
|
|
except duckdb.Error as e:
|
|
target.unlink(missing_ok=True)
|
|
raise ToolError(f"Die Datei ließ sich nicht als {fmt.upper()} lesen: {first_line(e)}") from e
|
|
if not columns:
|
|
target.unlink(missing_ok=True)
|
|
raise ToolError("Die Datei enthält keine Spalten.")
|
|
return columns, rows
|
|
|
|
|
|
def first_line(error: Exception) -> str:
|
|
return str(error).strip().splitlines()[0][:300] if str(error).strip() else type(error).__name__
|
|
|
|
|
|
@contextmanager
|
|
def readonly(tables: list[Table], settings: Settings) -> Iterator[duckdb.DuckDBPyConnection]:
|
|
con = duckdb.connect(":memory:", config={**BASE_CONFIG, "memory_limit": f"{settings.query_memory_mb}MB"})
|
|
try:
|
|
for i, table in enumerate(tables):
|
|
con.execute(f"ATTACH '{escape_literal(str(table.path))}' AS u{i} (READ_ONLY)")
|
|
con.execute(f'CREATE VIEW "{table.name}" AS SELECT * FROM u{i}.data')
|
|
con.execute("SET enable_external_access = false")
|
|
con.execute("SET lock_configuration = true")
|
|
yield con
|
|
finally:
|
|
con.close()
|
|
|
|
|
|
def escape_literal(value: str) -> str:
|
|
return value.replace("'", "''")
|
|
|
|
|
|
def query(tables: list[Table], sql: str, settings: Settings) -> tuple[list[str], list[tuple[Any, ...]], bool]:
|
|
"""Eine SELECT-Anweisung ausführen; liefert Spalten, höchstens query_max_rows Zeilen und ob gekürzt."""
|
|
with readonly(tables, settings) as con:
|
|
try:
|
|
statements = con.extract_statements(sql)
|
|
except duckdb.Error as e:
|
|
raise ToolError(f"SQL nicht lesbar: {first_line(e)}") from e
|
|
if len(statements) != 1 or statements[0].type != duckdb.StatementType.SELECT:
|
|
raise ToolError("Erlaubt ist genau eine SELECT-Anweisung (auch mit WITH), nichts anderes.")
|
|
timer = threading.Timer(settings.query_timeout_seconds, con.interrupt)
|
|
timer.start()
|
|
try:
|
|
result = con.execute(sql)
|
|
columns = [d[0] for d in result.description or []]
|
|
rows = result.fetchmany(settings.query_max_rows + 1)
|
|
except duckdb.InterruptException as e:
|
|
raise ToolError(f"Abfrage nach {settings.query_timeout_seconds:g} Sekunden abgebrochen; "
|
|
"bitte eingrenzen oder aggregieren.") from e
|
|
except duckdb.Error as e:
|
|
raise ToolError(f"SQL-Fehler: {first_line(e)}") from e
|
|
finally:
|
|
timer.cancel()
|
|
return columns, rows[:settings.query_max_rows], len(rows) > settings.query_max_rows
|
|
|
|
|
|
def as_table(columns: list[str], rows: list[tuple[Any, ...]]) -> str:
|
|
"""Ergebnis als Markdown-Tabelle für das Modell; lange Zellen gekürzt."""
|
|
def cell(value: Any) -> str:
|
|
text = "NULL" if value is None else str(value)
|
|
text = text.replace("\n", " ").replace("|", "\\|")
|
|
return text if len(text) <= CELL_CHARS else text[:CELL_CHARS] + "…"
|
|
lines = ["| " + " | ".join(columns) + " |", "|" + "---|" * len(columns)]
|
|
lines += ["| " + " | ".join(cell(v) for v in row) + " |" for row in rows]
|
|
return "\n".join(lines)
|