src/catalog/sample_decode.py

"""Base64-sample-value decoding for catalogs affected by a dedorch bug.

Go's introspection JSON-marshals numeric sample bytes as base64 (a `[]byte`
serialization quirk), so every decimal/int-typed column's `sample_values`
currently arrive as base64 strings (e.g. ``'OTUuMA=='`` for ``"95.0"``)
instead of the plain numeric text the planner LLM expects. The planner then
sees gibberish instead of value ranges for exactly the columns it filters and
aggregates on. Until Go fixes the marshaling, decode these at catalog read
time.

Conservative by design (a wrong decode silently corrupts planner context):
  - only numeric-typed columns are considered
  - every non-null sample in the column must pass a strict base64 gate
    (valid base64, decodes to printable ASCII, parses as a float) — a single
    non-conforming entry leaves the WHOLE column untouched (mixed content is
    suspicious, never guessed)
  - columns with `sample_values is None` (e.g. PII-flagged columns, which
    carry no samples by design) are skipped cleanly
  - self-disabling: once Go ships real numeric samples (plain ``"95.0"`` or
    actual numbers), the gate fails — plain digit strings are either not
    base64-padded correctly or don't decode to printable numeric text — so
    the pass becomes a no-op with no further changes needed here
"""

from __future__ import annotations

import base64
import binascii

from src.middlewares.logging import get_logger

from .models import Catalog

logger = get_logger("sample_decode")

_NUMERIC_TYPES = {
    "int",
    "integer",
    "bigint",
    "decimal",
    "numeric",
    "float",
    "double",
    "number",
}


def _decode_one(value: str) -> str | None:
    """Return the decoded numeric text for `value`, or None if it fails the gate."""
    if len(value) < 2 or len(value) % 4 != 0:
        return None
    try:
        decoded = base64.b64decode(value, validate=True)
    except (binascii.Error, ValueError):
        return None
    try:
        text = decoded.decode("ascii")
    except UnicodeDecodeError:
        return None
    if not text.isprintable():
        return None
    try:
        float(text)
    except ValueError:
        return None
    return text


def _decode_column_samples(samples: list) -> tuple[list, int] | None:
    """Return (decoded list, count decoded) if every non-null entry passes the gate.

    Returns None if any entry fails the gate (mixed content is left untouched).
    """
    decoded_values = []
    count = 0
    for entry in samples:
        if entry is None:
            decoded_values.append(None)
            continue
        if not isinstance(entry, str):
            return None
        decoded = _decode_one(entry)
        if decoded is None:
            return None
        count += 1
        decoded_values.append(decoded)
    if not count:
        return None
    return decoded_values, count


def decode_sample_values(catalog: Catalog) -> int:
    """Decode base64-encoded numeric sample_values in place. Returns count decoded.

    Never raises: any unexpected shape (wrong types, malformed entries) leaves
    the offending column's values untouched.
    """
    total = 0
    try:
        for source in catalog.sources:
            for table in source.tables:
                for col in table.columns:
                    if col.data_type.lower() not in _NUMERIC_TYPES:
                        continue
                    samples = col.sample_values
                    if not samples:
                        continue
                    result = _decode_column_samples(samples)
                    if result is None:
                        continue
                    decoded_values, count = result
                    col.sample_values = decoded_values
                    total += count
    except Exception as e:
        logger.error("sample decode failed", error=repr(e))
        return total
    if total:
        logger.info("decoded base64 sample values", user_id=catalog.user_id, count=total)
    return total