fix(dwh): align adapter sampling contract
This commit is contained in:
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user