Files

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")