stack/tests/aco/test_transpile_databricks.py

"""Integration tests — every pipeline transpiles to valid Databricks SQL.

For each of the 12 pipe modules, runs ``transpile()`` with a schema-only
DuckDB connection and verifies that every expression produces SQL that:

1. Is not an error comment
2. Parses as valid Databricks dialect via sqlglot
3. Contains the catalog prefix in all table references
"""

from __future__ import annotations

import importlib

import pytest
import sqlglot
from sqlglot import exp

from aco.lake.transpile import transpile

CATALOG = "test_catalog"

PIPE_MODULES = [
    "aco.pipe.core",
    "aco.pipe.readmissions",
    "aco.pipe.ahrq_measures",
    "aco.pipe.pharmacy",
    "aco.pipe.hcc_suspecting",
    "aco.pipe.input_layer",
    "aco.pipe.claims_preprocessing",
    "aco.pipe.cclf",
    "aco.pipe.data_quality",
    "aco.pipe.main",
    "aco.pipe.quality_measures",
    "aco.pipe.provider_attribution",
]

_IDS = [m.split(".")[-1] for m in PIPE_MODULES]


def _load_pipeline(module_path: str):
    mod = importlib.import_module(module_path)
    return mod.pipeline


# ── Per-module transpile results (cached) ─────────────────────────


@pytest.fixture(scope="module")
def all_transpiled():
    """Transpile every pipeline, returning {module: {expr_name: sql}}.

    Pipelines that hit sqlglot RecursionError (e.g. deeply nested
    DuckDB SQL from pivot expressions) are transpiled expression-by-
    expression so that passing expressions are still tested.
    """
    from aco.pipe.base import Pipeline

    results = {}
    for mod_path in PIPE_MODULES:
        short = mod_path.split(".")[-1]
        pipeline = _load_pipeline(mod_path)
        try:
            sql_map = transpile(
                pipeline,
                target_dialect="databricks",
                catalog=CATALOG,
                output_mode="select",
            )
        except RecursionError:
            # Fall back to per-expression transpilation so only the
            # offending expression is marked, not the whole module.
            sql_map = {}
            for i, expr in enumerate(pipeline.exprs):
                sub = Pipeline(exprs=pipeline.exprs[: i + 1])
                try:
                    sub_map = transpile(
                        sub,
                        target_dialect="databricks",
                        catalog=CATALOG,
                        output_mode="select",
                    )
                    sql_map[expr.name] = sub_map[expr.name]
                except RecursionError:
                    sql_map[expr.name] = (
                        "-- ERROR: sqlglot RecursionError "
                        "(DuckDB SQL too deeply nested)"
                    )
        results[short] = sql_map
    return results


def _flat_cases(all_transpiled):
    """Flatten {module: {name: sql}} to [(module, name, sql), ...]."""
    cases = []
    for mod, sql_map in all_transpiled.items():
        for name, sql in sql_map.items():
            cases.append((mod, name, sql))
    return cases


# ── Test: every expression transpiles without error ───────────────


# Schema-only transpilation limitation: nw.concat requires identical
# schemas, but the three medical_claim intermediate views produce
# different column counts in schema-only mode (DuckDB doesn't fully
# preserve dynamically-generated null columns from narwhals expressions).
# Works correctly with real data.
_KNOWN_TRANSPILE_FAILURES: set[str] = set()

_KNOWN_PARSE_FAILURES: set[str] = set()


class TestTranspileSuccess:
    """Every expression must transpile without a ``-- ERROR`` comment."""

    @pytest.mark.parametrize("mod_path", PIPE_MODULES, ids=_IDS)
    def test_no_errors(self, mod_path, all_transpiled):
        short = mod_path.split(".")[-1]
        sql_map = all_transpiled[short]
        errors = {
            name: sql
            for name, sql in sql_map.items()
            if sql.startswith("-- ERROR") and name not in _KNOWN_TRANSPILE_FAILURES
        }
        assert not errors, (
            f"{short}: {len(errors)} expression(s) failed to transpile:\n"
            + "\n".join(f"  {n}: {s}" for n, s in errors.items())
        )

    @pytest.mark.parametrize("mod_path", PIPE_MODULES, ids=_IDS)
    def test_known_recursion_issues_are_documented(self, mod_path, all_transpiled):
        """Track known sqlglot recursion issues — fail if fixed upstream."""
        short = mod_path.split(".")[-1]
        sql_map = all_transpiled[short]
        fixed = {
            name
            for name in _KNOWN_TRANSPILE_FAILURES
            if name in sql_map and not sql_map[name].startswith("-- ERROR")
        }
        if fixed:
            pytest.fail(
                f"These expressions are no longer broken — remove from "
                f"_KNOWN_TRANSPILE_FAILURES: {fixed}"
            )


# ── Test: every SQL parses as valid Databricks ────────────────────


class TestValidDatabricksSQL:
    """Every transpiled query must parse as valid Databricks SQL."""

    @pytest.mark.parametrize("mod_path", PIPE_MODULES, ids=_IDS)
    def test_all_parse_as_databricks(self, mod_path, all_transpiled):
        short = mod_path.split(".")[-1]
        sql_map = all_transpiled[short]
        failures = {}
        for name, sql in sql_map.items():
            if sql.startswith("-- ERROR"):
                continue
            if name in _KNOWN_PARSE_FAILURES:
                continue
            try:
                tree = sqlglot.parse_one(sql, dialect="databricks")
                assert isinstance(tree, (exp.Select, exp.Union))
            except Exception as exc:
                failures[name] = str(exc)
        assert not failures, (
            f"{short}: {len(failures)} expression(s) not valid Databricks SQL:\n"
            + "\n".join(f"  {n}: {e}" for n, e in failures.items())
        )


# ── Test: all table refs carry the catalog prefix ─────────────────


class TestCatalogReferences:
    """All table references must use 3-level catalog.schema.table."""

    @pytest.mark.parametrize("mod_path", PIPE_MODULES, ids=_IDS)
    def test_all_tables_have_catalog(self, mod_path, all_transpiled):
        short = mod_path.split(".")[-1]
        sql_map = all_transpiled[short]
        missing = []
        for name, sql in sql_map.items():
            if sql.startswith("-- ERROR"):
                continue
            try:
                tree = sqlglot.parse_one(sql, dialect="databricks")
            except Exception:
                continue
            for table in tree.find_all(exp.Table):
                if not table.name or not table.db:
                    continue
                if not table.catalog:
                    missing.append(f"{name}: {table.db}.{table.name} missing catalog")
        assert not missing, f"{short}: table refs without catalog:\n" + "\n".join(
            f"  {m}" for m in missing
        )

    @pytest.mark.parametrize("mod_path", PIPE_MODULES, ids=_IDS)
    def test_catalog_value_is_correct(self, mod_path, all_transpiled):
        short = mod_path.split(".")[-1]
        sql_map = all_transpiled[short]
        wrong = []
        for name, sql in sql_map.items():
            if sql.startswith("-- ERROR"):
                continue
            try:
                tree = sqlglot.parse_one(sql, dialect="databricks")
            except Exception:
                continue
            for table in tree.find_all(exp.Table):
                if table.catalog and table.catalog != CATALOG:
                    wrong.append(f"{name}: expected {CATALOG!r}, got {table.catalog!r}")
        assert not wrong, f"{short}: wrong catalog values:\n" + "\n".join(
            f"  {w}" for w in wrong
        )


# ── Test: output modes wrap SQL correctly ─────────────────────────


class TestOutputModes:
    """Verify insert, ctas, and view modes produce correct wrappers."""

    @pytest.fixture(scope="class")
    def core_insert(self):
        return transpile(
            _load_pipeline("aco.pipe.core"),
            target_dialect="databricks",
            catalog=CATALOG,
            output_mode="insert",
        )

    @pytest.fixture(scope="class")
    def core_ctas(self):
        return transpile(
            _load_pipeline("aco.pipe.core"),
            target_dialect="databricks",
            catalog=CATALOG,
            output_mode="ctas",
        )

    @pytest.fixture(scope="class")
    def core_view(self):
        return transpile(
            _load_pipeline("aco.pipe.core"),
            target_dialect="databricks",
            catalog=CATALOG,
            output_mode="view",
        )

    def test_insert_starts_correctly(self, core_insert):
        for name, sql in core_insert.items():
            if sql.startswith("-- ERROR"):
                continue
            assert sql.startswith("INSERT INTO"), (
                f"{name}: expected INSERT INTO, got {sql[:40]!r}"
            )

    def test_insert_target_has_catalog(self, core_insert):
        for name, sql in core_insert.items():
            if sql.startswith("-- ERROR"):
                continue
            assert f"INSERT INTO {CATALOG}." in sql, (
                f"{name}: INSERT target missing catalog"
            )

    def test_ctas_starts_correctly(self, core_ctas):
        for name, sql in core_ctas.items():
            if sql.startswith("-- ERROR"):
                continue
            assert sql.startswith("CREATE OR REPLACE TABLE"), (
                f"{name}: expected CTAS, got {sql[:50]!r}"
            )

    def test_view_starts_correctly(self, core_view):
        for name, sql in core_view.items():
            if sql.startswith("-- ERROR"):
                continue
            assert sql.startswith("CREATE OR REPLACE VIEW"), (
                f"{name}: expected VIEW, got {sql[:50]!r}"
            )


# ── Test: no temp view artifacts leak into output ─────────────────


class TestNoTempViewLeaks:
    """Transpiled SQL must not contain DuckDB temp view names."""

    @pytest.mark.parametrize("mod_path", PIPE_MODULES, ids=_IDS)
    def test_no_tv_prefix_in_sql(self, mod_path, all_transpiled):
        short = mod_path.split(".")[-1]
        sql_map = all_transpiled[short]
        leaks = []
        for name, sql in sql_map.items():
            if sql.startswith("-- ERROR"):
                continue
            if "_tv_" in sql:
                leaks.append(name)
        assert not leaks, f"{short}: temp view names leaked in: {leaks}"

    @pytest.mark.parametrize("mod_path", PIPE_MODULES, ids=_IDS)
    def test_no_unnamed_relation_in_sql(self, mod_path, all_transpiled):
        short = mod_path.split(".")[-1]
        sql_map = all_transpiled[short]
        leaks = []
        for name, sql in sql_map.items():
            if sql.startswith("-- ERROR"):
                continue
            if "unnamed_relation" in sql:
                leaks.append(name)
        assert not leaks, f"{short}: unnamed_relation artifacts in: {leaks}"