158 lines
7.5 KiB
Python
158 lines
7.5 KiB
Python
"""Uploads mit DuckDB-Analyse (→ Konzept 7.7)."""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from sqlalchemy import select
|
|
|
|
from app.core.crypto import Vault
|
|
from app.data_sources.connectors.git_repo import RepoCache
|
|
from app.identity import Principal
|
|
from app.tasks.worker import Worker
|
|
from app.tools import service as tools
|
|
from app.tools.context import ToolContext
|
|
from app.tools.registry import ToolError, load
|
|
from app.uploads import duck
|
|
from app.uploads import service as uploads
|
|
from app.uploads.models import Upload
|
|
from helpers import create_workspace
|
|
|
|
CSV = b"region,umsatz,monat\nNord,100,2025-01\nSued,250,2025-01\nNord,50,2025-02\n"
|
|
|
|
|
|
@pytest.fixture
|
|
def ws(browser, people):
|
|
anna = browser("anna")
|
|
workspace_id = create_workspace(anna, people["team_demo"])
|
|
anna.post(f"/workspaces/{workspace_id}/members", {"username": "ben", "role": "member"})
|
|
return {"anna": anna, "ben": browser("ben"), "id": workspace_id}
|
|
|
|
|
|
def upload(b, workspace_id: int, name: str, content: bytes):
|
|
return b.post(f"/workspaces/{workspace_id}/uploads", files={"file": (name, content)})
|
|
|
|
|
|
def rows(db) -> list[Upload]:
|
|
with db.session() as session:
|
|
result = list(session.scalars(select(Upload).order_by(Upload.id)))
|
|
session.expunge_all()
|
|
return result
|
|
|
|
|
|
def test_hochladen_csv_json_parquet(ws, db, settings, tmp_path):
|
|
assert upload(ws["anna"], ws["id"], "Umsatz 2025.csv", CSV).status_code == 303
|
|
upload(ws["ben"], ws["id"], "events.jsonl", b'{"typ":"a","n":1}\n{"typ":"b","n":2}\n')
|
|
import duckdb
|
|
parquet = tmp_path / "p.parquet"
|
|
duckdb.sql("SELECT 1 AS x, 'y' AS z").write_parquet(str(parquet))
|
|
upload(ws["anna"], ws["id"], "Umsatz 2025.parquet", parquet.read_bytes())
|
|
found = rows(db)
|
|
assert [(u.table_name, u.status, u.row_count) for u in found] == [
|
|
("umsatz_2025", "ok", 3), ("events", "ok", 2), ("umsatz_2025_2", "ok", 1)]
|
|
assert found[0].columns[1] == {"name": "umsatz", "type": "BIGINT"}
|
|
assert Path(settings.upload_dir, str(ws["id"]), f"{found[0].id}.duckdb").exists()
|
|
page = ws["anna"].get(f"/workspaces/{ws['id']}/data").text
|
|
assert "umsatz_2025" in page and "3 Zeilen, 3 Spalten" in page
|
|
|
|
|
|
def test_ungueltige_dateien(ws, db, settings):
|
|
assert upload(ws["anna"], ws["id"], "skript.py", b"print(1)").status_code == 400
|
|
assert upload(ws["anna"], ws["id"], "leer.csv", b"").status_code == 400
|
|
upload(ws["anna"], ws["id"], "kaputt.json", b"{nicht json")
|
|
assert rows(db)[-1].status == "error" and "JSON" in rows(db)[-1].error
|
|
settings.upload_max_mb = 0
|
|
assert upload(ws["anna"], ws["id"], "gross.csv", CSV).status_code == 400
|
|
# Fremde sehen den Workspace nicht
|
|
assert upload(ws["anna"].__class__(ws["anna"].client), ws["id"], "x.csv", CSV).status_code in (303, 403, 404)
|
|
|
|
|
|
def test_loeschen_nur_eigene_oder_admin(ws, db, settings, browser):
|
|
upload(ws["anna"], ws["id"], "a.csv", CSV)
|
|
upload(ws["ben"], ws["id"], "b.csv", CSV)
|
|
a, b = rows(db)
|
|
assert ws["ben"].post(f"/uploads/{a.id}/delete").status_code == 403
|
|
assert browser("dora").post(f"/uploads/{b.id}/delete").status_code == 404
|
|
assert ws["anna"].post(f"/uploads/{b.id}/delete").status_code == 303
|
|
assert ws["ben"].post(f"/uploads/{b.id}/delete").status_code == 404
|
|
assert not Path(settings.upload_dir, str(ws["id"]), f"{b.id}.duckdb").exists()
|
|
|
|
|
|
def test_chat_mit_upload_ohne_git(ws, db):
|
|
response = ws["anna"].post(f"/workspaces/{ws['id']}/chats")
|
|
assert response.status_code == 400
|
|
upload(ws["anna"], ws["id"], "a.csv", CSV)
|
|
assert ws["anna"].post(f"/workspaces/{ws['id']}/chats").status_code == 303
|
|
|
|
|
|
def tables(db, settings, workspace_id):
|
|
with db.session() as session:
|
|
return uploads.tables(session, settings, workspace_id)
|
|
|
|
|
|
def test_abfrage_nur_lesend(ws, db, settings, people):
|
|
upload(ws["anna"], ws["id"], "umsatz.csv", CSV)
|
|
other = create_workspace(ws["anna"], people["team_demo"], "Anderer")
|
|
upload(ws["anna"], other, "geheim.csv", CSV)
|
|
t = tables(db, settings, ws["id"])
|
|
foreign = tables(db, settings, other)[0].path
|
|
columns, result, truncated = duck.query(
|
|
t, "SELECT region, sum(umsatz) AS s FROM umsatz GROUP BY region ORDER BY region", settings)
|
|
assert columns == ["region", "s"] and result == [("Nord", 150), ("Sued", 250)] and not truncated
|
|
for sql in ["DROP TABLE umsatz", "SELECT 1; SELECT 2", "INSERT INTO umsatz VALUES ('x', 1, 'y')",
|
|
"COPY umsatz TO 'x.csv'", "ATTACH 'x.duckdb' AS x", "SET enable_external_access = true",
|
|
"INSTALL httpfs", "PRAGMA version"]:
|
|
with pytest.raises(ToolError):
|
|
duck.query(t, sql, settings)
|
|
for sql in ["SELECT * FROM read_csv('/etc/passwd')", "SELECT * FROM read_text('/etc/hosts')",
|
|
f"SELECT * FROM '{foreign}'", f"SELECT * FROM read_csv('{foreign}')", "SELECT * FROM read_parquet('https://example.org/x.parquet')"]:
|
|
with pytest.raises(ToolError, match="SQL-Fehler"):
|
|
duck.query(t, sql, settings)
|
|
|
|
|
|
def test_abfrage_limits(ws, db, settings):
|
|
upload(ws["anna"], ws["id"], "umsatz.csv", CSV)
|
|
t = tables(db, settings, ws["id"])
|
|
settings.query_max_rows = 2
|
|
_, result, truncated = duck.query(t, "SELECT * FROM umsatz", settings)
|
|
assert len(result) == 2 and truncated
|
|
settings.query_timeout_seconds = 0.2
|
|
with pytest.raises(ToolError, match="abgebrochen"):
|
|
duck.query(t, "SELECT count(*) FROM range(10000000000) a", settings)
|
|
|
|
|
|
def test_tools_und_chat(ws, db, settings, people, fake_provider):
|
|
upload(ws["anna"], ws["id"], "umsatz.csv", CSV)
|
|
admin = Principal(user_id=people["anna"], username="anna", display_name="Anna", is_admin=True,
|
|
team_ids=frozenset({people["team_demo"]}))
|
|
with db.session() as session:
|
|
tools.classify_suggested(session, admin)
|
|
assert {"data_list_tables", "data_describe_table", "data_query"} <= set(load())
|
|
chat_id = int(ws["anna"].post(f"/workspaces/{ws['id']}/chats").headers["location"].rsplit("/", 1)[1])
|
|
fake_provider.script = [
|
|
[("data_list_tables", {}), ("data_describe_table", {"table": "umsatz"})],
|
|
[("data_query", {"sql": "SELECT region, sum(umsatz) AS s FROM umsatz GROUP BY 1 ORDER BY 1"}),
|
|
("data_query", {"sql": "DELETE FROM umsatz"})],
|
|
"Nord 150, Süd 250.",
|
|
]
|
|
ws["anna"].post(f"/chats/{chat_id}/messages", {"content": "Umsatz je Region?"}, headers={"HX-Request": "true"})
|
|
Worker(db, settings).run_once()
|
|
calls = fake_provider.calls
|
|
assert {t["name"] for t in calls[0]["tools"]} == {"data_list_tables", "data_describe_table", "data_query"}
|
|
assert "Hochgeladene Dateien" in calls[0]["instructions"]
|
|
outputs = [i["output"] for i in calls[1]["input"] if i.get("type") == "function_call_output"]
|
|
assert "umsatz: 3 Zeilen; Spalten: region (VARCHAR), umsatz (BIGINT)" in outputs[0]
|
|
assert "| region | umsatz | monat |" in outputs[1]
|
|
outputs = [i["output"] for i in calls[2]["input"] if i.get("type") == "function_call_output"][2:]
|
|
assert "| Nord | 150 |" in outputs[0] and "2 Zeilen." in outputs[0]
|
|
assert "genau eine SELECT" in outputs[1]
|
|
assert "Nord 150" in ws["anna"].get(f"/chats/{chat_id}").text
|
|
|
|
|
|
def test_tools_ohne_uploads(settings):
|
|
context = ToolContext(settings=settings, vault=Vault(settings.secret_key), workspace_id=1, sources=[],
|
|
cache=RepoCache(settings))
|
|
assert "uploads" not in context.kinds
|
|
with pytest.raises(ToolError, match="keine hochgeladenen"):
|
|
load()["data_query"].run(context, load()["data_query"].validate({"sql": "SELECT 1"}))
|
|
json.dumps(load()["data_list_tables"].schema)
|