src/pipeline/db_pipeline/extractor.py

"""Schema introspection and per-column profiling for a user's database.

Identifiers (table/column names) are quoted via the engine's dialect preparer,
which handles reserved words, mixed case, and embedded quotes correctly across
dialects. Values used in SQL come from SQLAlchemy inspection of the DB itself,
not user input.
"""

from typing import Optional

import pandas as pd
from sqlalchemy import Float, Integer, Numeric, inspect
from sqlalchemy.engine import Engine

from src.middlewares.logging import get_logger

logger = get_logger("db_extractor")

TOP_VALUES_THRESHOLD = 0.05  # show top values if distinct_ratio <= 5%

# Dialects where PERCENTILE_CONT(...) WITHIN GROUP is supported as an aggregate.
# MySQL has no percentile aggregate; BigQuery has PERCENTILE_CONT only as an
# analytic (window) function — both drop median and keep min/max/mean.
_MEDIAN_DIALECTS = frozenset({"postgresql", "mssql", "snowflake"})


def _supports_median(engine: Engine) -> bool:
    return engine.dialect.name in _MEDIAN_DIALECTS


def _head_query(
    engine: Engine,
    select_clause: str,
    from_clause: str,
    n: int,
    order_by: str = "",
) -> str:
    """LIMIT/TOP-equivalent head query for the engine's dialect."""
    if engine.dialect.name == "mssql":
        return f"SELECT TOP {n} {select_clause} FROM {from_clause} {order_by}".strip()
    return f"SELECT {select_clause} FROM {from_clause} {order_by} LIMIT {n}".strip()


def _qi(engine: Engine, name: str) -> str:
    """Dialect-correct identifier quoting (schema.table also handled if dotted)."""
    preparer = engine.dialect.identifier_preparer
    if "." in name:
        schema, _, table = name.partition(".")
        return f"{preparer.quote(schema)}.{preparer.quote(table)}"
    return preparer.quote(name)


def get_schema(
    engine: Engine, exclude_tables: Optional[frozenset[str]] = None
) -> dict[str, list[dict]]:
    """Returns {table_name: [{name, type, is_numeric, is_primary_key, foreign_key}, ...]}."""
    exclude = exclude_tables or frozenset()
    inspector = inspect(engine)
    schema = {}
    for table_name in inspector.get_table_names():
        if table_name in exclude:
            continue

        pk = inspector.get_pk_constraint(table_name)
        pk_cols = set(pk["constrained_columns"]) if pk else set()

        fk_map = {}
        for fk in inspector.get_foreign_keys(table_name):
            for col, ref_col in zip(fk["constrained_columns"], fk["referred_columns"]):
                fk_map[col] = f"{fk['referred_table']}.{ref_col}"

        cols = inspector.get_columns(table_name)
        schema[table_name] = [
            {
                "name": c["name"],
                "type": str(c["type"]),
                "is_numeric": isinstance(c["type"], (Integer, Numeric, Float)),
                "is_primary_key": c["name"] in pk_cols,
                "foreign_key": fk_map.get(c["name"]),
            }
            for c in cols
        ]
    logger.info("extracted schema", table_count=len(schema))
    return schema


def get_row_count(engine: Engine, table_name: str) -> int:
    return pd.read_sql(f"SELECT COUNT(*) FROM {_qi(engine, table_name)}", engine).iloc[0, 0]


def profile_column(
    engine: Engine,
    table_name: str,
    col_name: str,
    is_numeric: bool,
    row_count: int,
) -> dict:
    """Returns null_count, distinct_count, min/max, top values, and sample values."""
    if row_count == 0:
        return {
            "null_count": 0,
            "distinct_count": 0,
            "distinct_ratio": 0.0,
            "sample_values": [],
        }

    qt = _qi(engine, table_name)
    qc = _qi(engine, col_name)

    # Combined stats query: null_count, distinct_count, and min/max (if numeric).
    # One round-trip instead of two.
    select_cols = [
        f"COUNT(*) - COUNT({qc}) AS nulls",
        f"COUNT(DISTINCT {qc}) AS distincts",
    ]
    if is_numeric:
        select_cols.append(f"MIN({qc}) AS min_val")
        select_cols.append(f"MAX({qc}) AS max_val")
        select_cols.append(f"AVG({qc}) AS mean_val")
        if _supports_median(engine):
            select_cols.append(
                f"PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY {qc}) AS median_val"
            )
    stats = pd.read_sql(f"SELECT {', '.join(select_cols)} FROM {qt}", engine)

    null_count = int(stats.iloc[0]["nulls"])
    distinct_count = int(stats.iloc[0]["distincts"])
    distinct_ratio = distinct_count / row_count if row_count > 0 else 0

    profile = {
        "null_count": null_count,
        "distinct_count": distinct_count,
        "distinct_ratio": round(distinct_ratio, 4),
    }

    if is_numeric:
        profile["min"] = stats.iloc[0]["min_val"]
        profile["max"] = stats.iloc[0]["max_val"]
        profile["mean"] = stats.iloc[0]["mean_val"]
        if _supports_median(engine):
            profile["median"] = stats.iloc[0]["median_val"]

    if 0 < distinct_ratio <= TOP_VALUES_THRESHOLD:
        top_sql = _head_query(
            engine,
            select_clause=f"{qc}, COUNT(*) AS cnt",
            from_clause=f"{qt} GROUP BY {qc}",
            n=10,
            order_by="ORDER BY cnt DESC",
        )
        top = pd.read_sql(top_sql, engine)
        profile["top_values"] = list(zip(top.iloc[:, 0].tolist(), top["cnt"].tolist()))

    sample = pd.read_sql(_head_query(engine, qc, qt, 5), engine)
    profile["sample_values"] = sample.iloc[:, 0].tolist()

    return profile


def profile_table(engine: Engine, table_name: str, columns: list[dict]) -> list[dict]:
    """Profile every column in a table. Returns [{col, profile, text}, ...].

    Per-column errors are logged and skipped so one bad column doesn't abort
    the whole table.
    """
    row_count = get_row_count(engine, table_name)
    if row_count == 0:
        logger.info("skipping empty table", table=table_name)
        return []

    results = []
    for col in columns:
        try:
            profile = profile_column(
                engine, table_name, col["name"], col.get("is_numeric", False), row_count
            )
            text = build_text(table_name, row_count, col, profile)
            results.append({"col": col, "profile": profile, "text": text})
        except Exception as e:
            logger.error(
                "column profiling failed",
                table=table_name,
                column=col["name"],
                error=str(e),
            )
            continue
    return results


def build_text(table_name: str, row_count: int, col: dict, profile: dict) -> str:
    col_name = col["name"]
    col_type = col["type"]

    key_label = ""
    if col.get("is_primary_key"):
        key_label = " [PRIMARY KEY]"
    elif col.get("foreign_key"):
        key_label = f" [FK -> {col['foreign_key']}]"

    text = f"Table: {table_name} ({row_count} rows)\n"
    text += f"Column: {col_name} ({col_type}){key_label}\n"
    text += f"Null count: {profile['null_count']}\n"
    text += f"Distinct count: {profile['distinct_count']} ({profile['distinct_ratio']:.1%})\n"
    if "min" in profile:
        text += f"Min: {profile['min']}, Max: {profile['max']}\n"
        text += f"Mean: {profile['mean']}\n"
        if profile.get("median") is not None:
            text += f"Median: {profile['median']}\n"
    if "top_values" in profile:
        top_str = ", ".join(f"{v} ({c})" for v, c in profile["top_values"])
        text += f"Top values: {top_str}\n"
    text += f"Sample values: {profile['sample_values']}"
    return text