43 lines
1.4 KiB
Python
43 lines
1.4 KiB
Python
"""PostgreSQL execution operations used by the direct DWH adapter."""
|
|
|
|
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:
|
|
return run_controlled(
|
|
engine,
|
|
sql,
|
|
limit=DEFAULT_LIMIT if limit is None else 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]
|