Files
ThothII/harness/tht/ctetest.py

125 lines
4.3 KiB
Python

import json
from datetime import datetime
from pathlib import Path
from typing import Literal
import sqlglot
from pydantic import BaseModel
from sqlglot import exp
CTE_TESTS_FILE = "cte_tests.json"
PROBE = "\nSELECT * FROM __tht_probe__"
class CteError(Exception):
pass
class CteTestRecord(BaseModel):
name: str
ts: datetime
sql_hash: str
status: Literal["ok", "error"]
columns: list[str] = []
row_sample: int = 0
execution_ms: int = 0
warnings: list[str] = []
error: str | None = None
preview_rows: list[list[str | int | float | bool | None]] | None = None
_MAX_CELL_CHARS = 200
def _jsonable(v):
"""Coerce a raw DB cell to a JSON-native value for persistence in cte_tests.json.
None/bool/int/float pass through unchanged; str (and anything else, via str())
is truncated to ~200 chars."""
if v is None or isinstance(v, (bool, int, float)):
return v
s = v if isinstance(v, str) else str(v)
return s[:_MAX_CELL_CHARS]
def has_trailing_select(cte_sql: str) -> bool:
"""True se il file, parsato da solo, e' gia' una query completa con SELECT
finale dopo il blocco WITH (violazione della regola: il file deve contenere
solo WITH ... AS (...)).
I file CTE corretti contengono solo 'WITH ... AS (...)': sqlglot.parse_one()
fallisce su di loro perche' manca la SELECT obbligatoria.
Un file con SELECT finale parsa invece correttamente come Select con un nodo With.
I file con errori di sintassi falliscono parse_one e ritornano False (no hint)."""
try:
ast = sqlglot.parse_one(cte_sql, read="postgres")
except sqlglot.errors.ParseError:
return False
return ast is not None and ast.find(exp.With) is not None
def last_cte_name(cte_sql: str) -> str:
"""Nome dell'ultimo CTE di un blocco 'WITH a AS (...), b AS (...)' senza SELECT
finale. Trucco AV-SQL: si appende una SELECT fittizia per rendere parsabile."""
if not cte_sql.lstrip().upper().startswith("WITH"):
raise CteError("il file deve contenere un blocco WITH ... AS (...)")
try:
ast = sqlglot.parse_one(cte_sql + PROBE, read="postgres")
except sqlglot.errors.ParseError as e:
raise CteError(f"CTE non parsabile: {e}") from e
with_clause = ast.find(exp.With)
if with_clause is None or not with_clause.expressions:
raise CteError("il file deve contenere un blocco WITH ... AS (...)")
return with_clause.expressions[-1].alias
def build_test_sql(cte_sql: str) -> str:
"""SQL completo per il test: il blocco WITH + SELECT * dall'ultimo CTE.
Il LIMIT lo inietta run_controlled."""
name = last_cte_name(cte_sql)
return f"{cte_sql.rstrip().rstrip(',')}\nSELECT * FROM {name}"
def _iter_json_objects(text: str):
"""Estrae gli oggetti JSON di primo livello da `text`, qualunque sia il
formato: array JSON, JSONL, o misto (array + righe JSONL appese). Salta
spazi, virgole e parentesi d'array fra un oggetto e il successivo."""
dec = json.JSONDecoder()
i, n = 0, len(text)
while i < n:
if text[i] in " \t\r\n,[]":
i += 1
continue
obj, end = dec.raw_decode(text, i)
yield obj
i = end
def load_cte_tests(session_dir: Path) -> list[CteTestRecord]:
path = session_dir / CTE_TESTS_FILE
if not path.exists():
return []
return load_cte_tests_text(path.read_text())
def load_cte_tests_text(text: str) -> list[CteTestRecord]:
try:
objs = list(_iter_json_objects(text))
except json.JSONDecodeError as e:
raise CteError(f"{CTE_TESTS_FILE} corrotto: {e}") from e
return [CteTestRecord.model_validate(o) for o in objs]
def append_cte_test_snapshot(repository, snapshot, record: CteTestRecord) -> None:
previous = snapshot.artifacts.get("cte_tests", "")
repository.write_artifact(snapshot.manifest.id, "cte_tests", previous + record.model_dump_json() + "\n")
def append_cte_test(session_dir: Path, record: CteTestRecord) -> None:
"""Append atomico (una riga JSON per esito): scritture concorrenti sulla
stessa sessione non si sovrascrivono, a differenza del rewrite dell'array."""
path = session_dir / CTE_TESTS_FILE
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("a", encoding="utf-8") as f:
f.write(record.model_dump_json() + "\n")