mptyl/ThothII
1
0
Fork 0
Code Issues 25 Pull Requests Actions Packages Projects Releases 1 Wiki Activity

fix(dwh): align adapter sampling contract

Browse Source
This commit is contained in:
marcopan
2026-07-11 20:14:23 +02:00
parent 717e5ecced
commit 216984aac8
11 changed files with 239 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
harness/tht/adapters/dwh/postgres.py
+11 -7
View File
@@ -1,12 +1,14 @@
"""Direct PostgreSQL implementation of the DWH port."""
from tht.config import DatabaseConfig
from tht.db import execute
from sqlalchemy.exc import SQLAlchemyError
from tht.db import execute, sampling
from tht.db.connection import make_engine, ping
from tht.db.introspect import introspect
from tht.execute import ExecResult, PlanSummary
from tht.mschema.models import PhysicalSchema
from tht.ports.dwh import DwhCapabilities, DwhHealth
from tht.ports.dwh import DistinctValues, DwhCapabilities, DwhHealth
class PostgresDwhAdapter:
@@ -19,23 +21,25 @@ class PostgresDwhAdapter:
def health(self) -> DwhHealth:
try:
ping(self._engine)
except Exception as exc:
except SQLAlchemyError as exc:
return DwhHealth(ok=False, detail=str(exc))
return DwhHealth(ok=True)
def introspect(self) -> PhysicalSchema:
return introspect(self._engine, self._config.database, self._config.db_schema)
def run_query(self, sql: str, *, limit: int | None = None) -> ExecResult:
def run_query(self, sql: str, *, limit: int) -> ExecResult:
return execute.run_query(self._engine, sql, limit=limit)
def explain(self, sql: str) -> PlanSummary:
return execute.explain(self._engine, sql)
def sample_column(self, table: str, column: str, *, limit: int) -> list[object]:
return execute.sample_column(
return sampling.sample_column(
self._engine, self._config.db_schema, table, column, limit=limit
)
def distinct_values(self, table: str, column: str) -> list[object]:
return execute.distinct_values(self._engine, self._config.db_schema, table, column)
def distinct_values(self, table: str, column: str) -> DistinctValues:
return sampling.distinct_values(
self._engine, self._config.db_schema, table, column
)
harness/tht/adapters/dwh/thoth_rest.py
+19 -17
View File
@@ -2,16 +2,12 @@
from tht.config import DatabaseConfig, RestConfig
from tht.db.introspect import introspect_rest
from tht.execute import ExecResult, PlanSummary
from tht.db import sampling
from tht.execute import ExecResult, ExecutionError, PlanSummary
from tht.mschema.models import PhysicalSchema
from tht.ports.dwh import DwhCapabilities, DwhHealth
from tht.ports.dwh import DistinctValues, DwhCapabilities, DwhHealth
from tht.rest.client import RestClient, RestError
from tht.rest.execute import (
distinct_values_rest,
explain_rest,
run_controlled_rest,
sample_column_rest,
)
from tht.rest.execute import explain_rest, run_controlled_rest
class ThothRestDwhAdapter:
@@ -34,18 +30,24 @@ class ThothRestDwhAdapter:
self._client, self._database.database, self._database.db_schema
)
def run_query(self, sql: str, *, limit: int | None = None) -> ExecResult:
return run_controlled_rest(self._client, sql, limit=10 if limit is None else limit)
def run_query(self, sql: str, *, limit: int) -> ExecResult:
return run_controlled_rest(self._client, sql, limit=limit)
def explain(self, sql: str) -> PlanSummary:
return explain_rest(self._client, sql)
def sample_column(self, table: str, column: str, *, limit: int) -> list[object]:
return sample_column_rest(
self._client, self._database.db_schema, table, column, limit=limit
)
try:
return sampling.sample_column_rest(
self._client, self._database.db_schema, table, column, limit=limit
)
except RestError as exc:
raise ExecutionError(str(exc)) from exc
def distinct_values(self, table: str, column: str) -> list[object]:
return distinct_values_rest(
self._client, self._database.db_schema, table, column
)
def distinct_values(self, table: str, column: str) -> DistinctValues:
try:
return sampling.distinct_values_rest(
self._client, self._database.db_schema, table, column
)
except RestError as exc:
raise ExecutionError(str(exc)) from exc
harness/tht/db/execute.py
+4 -24
View File
@@ -4,39 +4,19 @@ from sqlalchemy import Engine
from tht.execute import ExecResult, PlanSummary, explain as _explain, run_controlled
DEFAULT_LIMIT = 10
DEFAULT_TIMEOUT_MS = 30_000
def run_query(engine: Engine, sql: str, *, limit: int | None = None) -> ExecResult:
def run_query(engine: Engine, sql: str, *, limit: int) -> ExecResult:
if limit <= 0:
raise ValueError("limit must be a positive integer")
return run_controlled(
engine,
sql,
limit=DEFAULT_LIMIT if limit is None else limit,
limit=limit,
timeout_ms=DEFAULT_TIMEOUT_MS,
)
def explain(engine: Engine, sql: str) -> PlanSummary:
return _explain(engine, sql, timeout_ms=DEFAULT_TIMEOUT_MS)
def sample_column(
engine: Engine, schema: str, table: str, column: str, *, limit: int
) -> list[object]:
quote = engine.dialect.identifier_preparer.quote
sql = (
f"SELECT {quote(column)} FROM {quote(schema)}.{quote(table)} "
f"WHERE {quote(column)} IS NOT NULL"
)
return [row[0] for row in run_query(engine, sql, limit=limit).rows]
def distinct_values(engine: Engine, schema: str, table: str, column: str) -> list[object]:
quote = engine.dialect.identifier_preparer.quote
sql = (
f"SELECT DISTINCT {quote(column)} FROM {quote(schema)}.{quote(table)} "
f"WHERE {quote(column)} IS NOT NULL ORDER BY {quote(column)}"
)
result = run_controlled(engine, sql, limit=1000, timeout_ms=DEFAULT_TIMEOUT_MS)
return [row[0] for row in result.rows]
harness/tht/db/sampling.py
+56
View File
@@ -5,10 +5,66 @@ from sqlalchemy import Engine, text
from tht.config import ExamplesConfig, LshConfig
from tht.mschema.models import Annotations, PhysicalSchema
from tht.ports.dwh import DistinctValues
logger = logging.getLogger(__name__)
TEXT_TYPE_PREFIXES = ("text", "varchar", "character", "char")
DEFAULT_DISTINCT_VALUES_LIMIT = 1000
def _quoted_top_values_query(engine: Engine, schema: str, table: str, column: str):
quote = engine.dialect.identifier_preparer.quote
identifier = quote(column)
return text(
f"SELECT {identifier} AS value FROM {quote(schema)}.{quote(table)} "
f"WHERE {identifier} IS NOT NULL GROUP BY {identifier} "
f"ORDER BY count(*) DESC, {identifier} LIMIT :lim"
)
def sample_column(
engine: Engine, schema: str, table: str, column: str, *, limit: int
) -> list[object]:
if limit <= 0:
raise ValueError("limit must be a positive integer")
query = _quoted_top_values_query(engine, schema, table, column)
with engine.connect() as conn:
rows = conn.execute(query, {"lim": limit}).fetchall()
return [row[0] for row in rows]
def sample_column_rest(
client, schema: str, table: str, column: str, *, limit: int
) -> list[object]:
if limit <= 0:
raise ValueError("limit must be a positive integer")
rows = client.top_values(schema, table, column, limit)
return [row["value"] for row in rows if row.get("value") is not None]
def distinct_values(
engine: Engine,
schema: str,
table: str,
column: str,
*,
max_values: int = DEFAULT_DISTINCT_VALUES_LIMIT,
) -> DistinctValues:
values = sample_column(engine, schema, table, column, limit=max_values + 1)
return DistinctValues(values=values[:max_values], truncated=len(values) > max_values)
def distinct_values_rest(
client,
schema: str,
table: str,
column: str,
*,
max_values: int = DEFAULT_DISTINCT_VALUES_LIMIT,
) -> DistinctValues:
values = sample_column_rest(client, schema, table, column, limit=max_values + 1)
return DistinctValues(values=values[:max_values], truncated=len(values) > max_values)
def is_text_type(pg_type: str) -> bool:
harness/tht/ports/__init__.py
+2
View File
@@ -4,6 +4,7 @@ from tht.ports.dwh import (
DwhAdapter,
DwhCapabilities,
DwhHealth,
DistinctValues,
UnsupportedCapability,
)
@@ -11,5 +12,6 @@ __all__ = [
"DwhAdapter",
"DwhCapabilities",
"DwhHealth",
"DistinctValues",
"UnsupportedCapability",
]
harness/tht/ports/dwh.py
+8 -2
View File
@@ -21,6 +21,12 @@ class DwhHealth:
detail: str | None = None
@dataclass(frozen=True)
class DistinctValues:
values: list[object]
truncated: bool
class UnsupportedCapability(Exception):
"""Raised when an adapter cannot provide an optional DWH operation."""
@@ -34,10 +40,10 @@ class DwhAdapter(Protocol):
def introspect(self) -> PhysicalSchema: ...
def run_query(self, sql: str, *, limit: int | None = None) -> ExecResult: ...
def run_query(self, sql: str, *, limit: int) -> ExecResult: ...
def explain(self, sql: str) -> PlanSummary: ...
def sample_column(self, table: str, column: str, *, limit: int) -> list[object]: ...
def distinct_values(self, table: str, column: str) -> list[object]: ...
def distinct_values(self, table: str, column: str) -> DistinctValues: ...
harness/tht/rest/execute.py
+2 -11
View File
@@ -17,6 +17,8 @@ from tht.rest.explain import parse_text_plan
def run_controlled_rest(client, sql: str, *, limit: int) -> ExecResult:
# Guard read-only client-side anche sul path REST (D7): non delegare l'unica verifica
# al server. Stesso check strutturale del path diretto.
if limit <= 0:
raise ValueError("limit must be a positive integer")
assert_read_only(sql)
final_sql, injected = _inject_limit(sql, limit)
start = time.monotonic()
@@ -46,14 +48,3 @@ def explain_rest(client, sql: str) -> PlanSummary:
except RestError as e:
raise ExecutionError(str(e)) from e
return parse_text_plan(lines)
def sample_column_rest(
client, schema: str, table: str, column: str, *, limit: int
) -> list[object]:
rows = client.top_values(schema, table, column, limit)
return [row["value"] for row in rows if row.get("value") is not None]
def distinct_values_rest(client, schema: str, table: str, column: str) -> list[object]:
return sample_column_rest(client, schema, table, column, limit=1000)