server/ingestion.py

"""CSV → SQLite ingestion in the trained DB-path layout (F003, Slice S1).

Robust-parse an (untrusted, messy) CSV with DuckDB ``read_csv_auto`` and land it
as a plain committed SQLite database at the exact path the *unchanged*
``SQLEnvironment`` already loads from — ``<root>/<db_id>/<db_id>.sqlite`` — so the
RL-trained agent loop queries the user's data with zero environment changes.

The one real ingestion gotcha is fixed here: an integer-coded column that
contains a NULL can be widened to ``float64`` so a naive ``to_sql`` stores
``3.0``/``REAL`` instead of ``3``/``INTEGER`` (corrupting coded-value semantics).
``_coerce_nullable_ints`` casts such columns back to pandas nullable ``Int64`` so
SQLite declares ``INTEGER`` affinity and ``WHERE col = 3`` matches.

This module is intentionally dependency-light: it imports only ``duckdb`` /
``pandas`` (the ``serve`` extra) plus the stdlib, and pulls NONE of the heavy
training deps (``trl``/``torch``/``transformers``) — pinned by a subprocess guard
in ``tests/unit/test_ingestion.py``.
"""

from dataclasses import dataclass
from pathlib import Path
import re
import sqlite3
import tempfile

import duckdb
import pandas as pd

try:
    from .sql_ident import is_valid_identifier
except ImportError:  # pragma: no cover - flat-layout / direct-run fallback
    from sql_ident import is_valid_identifier  # type: ignore[no-redef]

# db_id and SQL identifiers must match the env's ``^[A-Za-z0-9_]+$`` charset
# (the single source of that contract is ``server/sql_ident.py``).
_NON_WORD = re.compile(r"[^A-Za-z0-9]+")


@dataclass(frozen=True)
class IngestResult:
    """Lightweight result of an ingest. Internal bookkeeping; never serialized."""

    db_id: str
    db_path: Path
    table: str
    root: Path
    row_count: int
    column_mapping: dict[str, str]


def _normalize_db_id(name: str) -> str:
    """Normalize an arbitrary name to a non-empty ``^[A-Za-z0-9_]+$`` db_id.

    Lowercases, collapses runs of non-alphanumerics to a single ``_``, and strips
    leading/trailing ``_``. Mirrors the regex ``SQLEnvironment`` enforces so the
    derived id is one the env will accept. Raises ``ValueError`` when nothing
    usable survives (e.g. ``"---"`` or ``""``).
    """
    lowered = name.strip().lower()
    collapsed = _NON_WORD.sub("_", lowered)
    stripped = collapsed.strip("_")
    if not is_valid_identifier(stripped):
        raise ValueError(f"Could not derive a valid database id from '{name}'.")
    return stripped


def _normalize_headers(columns: list[str]) -> dict[str, str]:
    """Map each original header to a unique SQL-safe identifier (order-preserving).

    DuckDB owns blank/duplicate-header uniqueness: ``read_csv_auto`` renames blank
    headers to ``column0..`` and identical repeats to ``a_1``/``a_2`` BEFORE this
    runs, so inputs here are always non-blank and distinct.

    This helper sanitizes characters (non-word runs collapse to ``_``). Sanitizing
    can EMPTY a non-blank header that DuckDB passed through: an all-non-``[A-Za-z0-9]``
    header (e.g. ``"数量"`` or ``"%%%"``) collapses to ``""``, which would make
    ``to_sql`` raise ``ValueError: Empty table or column name specified``. So a
    positional ``col_<index>`` fallback is applied when the sanitized base is empty,
    keeping the ``^[A-Za-z0-9_]+$`` contract.

    It still de-dupes the one collision DuckDB does NOT prevent: two DISTINCT headers
    that collapse to the same safe name after sanitizing (e.g. ``"a b"`` and ``"a-b"``
    both → ``a_b``), suffixing ``_2``/``_3``. Keyed by the ORIGINAL header so the
    returned map round-trips back to the source columns.
    """
    mapping: dict[str, str] = {}
    used: set[str] = set()
    for index, original in enumerate(columns):
        base = _NON_WORD.sub("_", original.strip().lower()).strip("_")
        if not base:
            # Sanitizing emptied a non-blank, all-non-alphanumeric header (e.g.
            # ``"数量"``/``"%%%"`` DuckDB passes through): fall back to a positional
            # name so the column is never blank for ``to_sql``.
            base = f"col_{index}"
        candidate = base
        suffix = 2
        while candidate in used:
            candidate = f"{base}_{suffix}"
            suffix += 1
        used.add(candidate)
        mapping[original] = candidate
    return mapping


def _coerce_nullable_ints(df: pd.DataFrame) -> pd.DataFrame:
    """Cast integral ``float64`` columns to pandas nullable ``Int64`` (the fix).

    For each ``float64`` column whose non-null values are all integral, cast to
    ``Int64`` so the subsequent ``to_sql`` declares INTEGER affinity and stores
    ``3`` (not ``3.0``). Genuine floats, text/object columns, and already-integer
    columns are left untouched. Returns a new DataFrame.

    Note: under the current DuckDB/pandas, ``read_csv_auto().df()`` may already
    yield ``Int64`` for null-coded integer columns, so on this stack the helper is
    a safety net rather than the active fix. Its real, load-bearing effect — a
    genuine ``float64`` integral+NaN column would otherwise store ``REAL``/``3.0``
    — is proven directly by ``test_coerce_changes_sqlite_storage_to_integer``,
    which writes such a column to SQLite with and without the coercer.
    """
    result = df.copy()
    for column in result.columns:
        series = result[column]
        if not str(series.dtype).startswith("float"):
            continue
        non_null = series.dropna()
        if non_null.empty:
            # Fully-null float column: nothing integral to preserve; leave as-is.
            continue
        if (non_null == non_null.round()).all():
            result[column] = series.astype("Int64")
    return result


def _stem_from_source(source: str | Path | bytes) -> str:
    """Best-effort db_id seed from the source (filename stem, else a default)."""
    if isinstance(source, bytes):
        return "upload"
    return Path(source).stem


def _read_csv(source: str | Path | bytes) -> pd.DataFrame:
    """Parse a CSV path or raw bytes with DuckDB ``read_csv_auto`` → DataFrame."""
    con = duckdb.connect()
    try:
        if isinstance(source, bytes):
            with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as handle:
                handle.write(source)
                tmp_path = handle.name
            try:
                return con.execute("SELECT * FROM read_csv_auto(?)", [tmp_path]).df()
            finally:
                Path(tmp_path).unlink(missing_ok=True)
        path = Path(source)
        if not path.exists():
            raise FileNotFoundError(f"CSV not found: {path}.")
        return con.execute("SELECT * FROM read_csv_auto(?)", [str(path)]).df()
    finally:
        con.close()


def ingest_csv(
    source: str | Path | bytes,
    *,
    db_id: str | None = None,
    table: str = "data",
    root: str | Path = "data/uploads",
    if_exists: str = "error",
) -> IngestResult:
    """Parse a CSV and write a committed SQLite DB the unchanged ``SQLEnvironment``
    opens read-only at ``<root>/<db_id>/<db_id>.sqlite``.

    Args:
        source: CSV file path, or raw CSV ``bytes`` (e.g. an upload).
        db_id: Explicit db_id; if None, derived from the filename (or ``"upload"``
            for bytes) and normalized via ``_normalize_db_id``.
        table: SQL table name to write (default ``"data"``).
        root: Upload root dir (a runtime root, NOT ``data/databases/``). Created if
            missing.
        if_exists: Collision policy for an existing ``<db_id>`` dir:
            ``"error"`` raises; ``"replace"`` overwrites; ``"version"`` appends
            ``_2``/``_3`` to the db_id.

    Returns:
        ``IngestResult`` with the normalized db_id, absolute db_path, table, root,
        row_count and original→safe column mapping.

    Raises:
        ValueError: ``table`` is not ``^[A-Za-z0-9_]+$``, db_id cannot normalize
            to ``^[A-Za-z0-9_]+$``, the CSV is empty, or ``if_exists="error"`` and
            the target already exists.
        FileNotFoundError: ``source`` is a path that does not exist.
    """
    # Validate the table name against the SAME ``^[A-Za-z0-9_]+$`` contract
    # ``data_card._validate_table_name`` enforces downstream, so a name ingestion
    # would write is never one the data card later rejects (e.g. ``"my table"``).
    if not is_valid_identifier(table):
        raise ValueError(f"Invalid table name '{table}': must match ^[A-Za-z0-9_]+$.")

    # Validate the source path BEFORE any normalization/write so a missing file
    # never creates a directory.
    if not isinstance(source, bytes) and not Path(source).exists():
        raise FileNotFoundError(f"CSV not found: {Path(source)}.")

    df = _read_csv(source)
    if df.shape[1] == 0 or df.shape[0] == 0:
        raise ValueError("CSV appears to be empty.")

    # Normalize the db_id BEFORE touching the filesystem.
    resolved_id = _normalize_db_id(db_id or _stem_from_source(source))

    root_path = Path(root)
    resolved_id = _resolve_collision(root_path, resolved_id, if_exists)

    # Re-creating/rewriting the table invalidates any existing data-card sidecar:
    # its profile/schema would now be stale (wrong advisory context — the exact
    # ADR 0007/0009 failure). Drop it so a stale card is never served.
    if if_exists in {"replace", "version"}:
        _invalidate_data_card_sidecar(root_path, resolved_id)

    # Normalize headers + apply the Int64 fix.
    column_mapping = _normalize_headers(list(df.columns))
    df = df.rename(columns=column_mapping)
    df = _coerce_nullable_ints(df)

    target_dir = root_path / resolved_id
    db_path = (target_dir / f"{resolved_id}.sqlite").resolve()
    target_dir.mkdir(parents=True, exist_ok=True)

    connection = sqlite3.connect(db_path)
    try:
        df.to_sql(table, connection, index=False, if_exists="replace")
        connection.commit()
    finally:
        connection.close()

    return IngestResult(
        db_id=resolved_id,
        db_path=db_path,
        table=table,
        root=root_path,
        row_count=int(df.shape[0]),
        column_mapping=column_mapping,
    )


def _invalidate_data_card_sidecar(root: Path, db_id: str) -> None:
    """Unlink any existing ``<root>/<db_id>/<db_id>.datacard.json`` sidecar.

    Mirrors ``data_card._sidecar_path`` (kept inline to keep ingestion free of a
    data_card import). A no-op when no sidecar exists. Called on re-ingest so a
    stale card built against the OLD schema is never served after a replace.
    """
    sidecar = root / db_id / f"{db_id}.datacard.json"
    sidecar.unlink(missing_ok=True)


def _resolve_collision(root: Path, db_id: str, if_exists: str) -> str:
    """Apply the ``if_exists`` policy to a possibly-existing ``<root>/<db_id>``.

    ``"error"`` raises before any write; ``"replace"`` returns the id unchanged
    (the writer overwrites in place); ``"version"`` returns a de-duped id
    (``<db_id>_2``/``_3``) that still fullmatches ``^[A-Za-z0-9_]+$``.
    """
    target_dir = root / db_id
    if not target_dir.exists():
        return db_id
    if if_exists == "error":
        raise ValueError(f"Database '{db_id}' already exists.")
    if if_exists == "replace":
        return db_id
    if if_exists == "version":
        suffix = 2
        candidate = f"{db_id}_{suffix}"
        while (root / candidate).exists():
            suffix += 1
            candidate = f"{db_id}_{suffix}"
        return candidate
    raise ValueError(f"Unknown if_exists policy: '{if_exists}'.")