fix: upgrade Polyglot to 0.12.3 for PIVOT and UNPIVOT star lineage

Cargo.toml

@@ -15,7 +15,7 @@ fensu-policy = "=0.22.0"
 globset = "0.4"
 proc-macro2 = { version = "1", features = ["span-locations"] }
 pyo3 = { version = "0.26", features = ["abi3-py312"] }
-polyglot-sql = { git = "https://github.com/chio-labs/polyglot.git", rev = "dd7154035566933a0f9acae17a2c102afb4dbfde", default-features = false, features = [
+polyglot-sql = { git = "https://github.com/chio-labs/polyglot.git", rev = "f3774102f518f716cd648e4fdebb9d8558a0b63a", default-features = false, features = [
     "ast-tools",
     "builder",
     "dialect-bigquery",

pyproject.toml

@@ -15,7 +15,7 @@ dependencies = [
   "duckdb>=1.3.2",
   "filelock>=3.20.3",
   "orjson>=3.10.0",
-  "polyglot-sql-chio==0.12.2",
+  "polyglot-sql-chio==0.12.3",
   "pydantic-settings>=2,<3",
   "PyYAML>=6.0.2",
   "rich>=15.0.0",

tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py

@@ -231,3 +231,21 @@ class SetOperationTypeIntegrationTestCase:
 class SetOperationLineageIntegrationTestCase:
     description: str
     expected_sources: frozenset[tuple[str, str]]
+
+
+@dataclass(frozen=True)
+class ReshapedStarIntegrationTestCase:
+    """One PIVOT or UNPIVOT star query whose compiled columns must match DuckDB."""
+
+    description: str
+    query_sql: str
+    expected_columns: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class ReshapedStarLineageIntegrationTestCase:
+    """One PIVOT star query with the expected upstream columns of each output."""
+
+    description: str
+    query_sql: str
+    expected_sources: tuple[frozenset[tuple[str, str]], ...]

tests/integration/src/sqlbuild/compiler/pipeline/helpers.py

@@ -8,18 +8,30 @@
 from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter
 from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter
 from sqlbuild.compiler.compile.main._build_compile_inputs import build_compile_inputs
-from sqlbuild.compiler.compile.models import CompileAdapterContext, CompileAuditInput
+from sqlbuild.compiler.compile.models import (
+    CompileAdapterContext,
+    CompileAuditInput,
+    CompiledLineageColumnFact,
+    CompiledModel,
+    CompiledProject,
+)
 from sqlbuild.compiler.discovery.main.discover import discover_project_inputs
 from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs
 from sqlbuild.compiler.manifest.main.build import build_manifest
 from sqlbuild.compiler.pipeline.main.compile import run_compile_pipeline
+from sqlbuild.compiler.pipeline.main.project import compile_project
 from sqlbuild.compiler.pipeline.models import CompilePipelineOptions, CompilePipelineResult
 from sqlbuild.runtime.contracts.models import ConnectionHooks
 from sqlbuild.sql_values.types import CollectionRendering
 
 _SCHEMA_FIXTURE_PATH: Path = (
     Path(__file__).resolve().parents[5] / "fixtures" / "dbt_manifest_v12_schema.json"
 )
+_RESHAPED_STAR_PROJECT_TOML: str = 'name = "orders"\nadapter = "duckdb"\n[rules]\nselect = []\n'
+RESHAPED_STAR_UPSTREAM_MODELS: dict[str, str] = {
+    "staged_orders": "SELECT 1 AS customer_id, 'books' AS category, 10 AS amount",
+    "wide_orders": "SELECT 1 AS customer_id, 2 AS books, 3 AS games",
+}
 _SEMANTIC_BINDING_PROJECT_TOML: str = 'name = "semantic_binding"\nadapter = "duckdb"\n'
 _AUDIT_FACTORY_ADAPTER_CONTEXT: CompileAdapterContext = CompileAdapterContext(
     value_renderer=DuckDbAdapter(),
@@ -125,3 +137,34 @@ def write_semantic_binding_project(
     models_dir.mkdir()
     _ = (models_dir / "upstream.sql").write_text(upstream_sql, encoding="utf-8")
     _ = (models_dir / "downstream.sql").write_text(downstream_sql, encoding="utf-8")
+
+
+def compile_reshaped_star_model(
+    *, project_dir: Path, query_sql: str
+) -> tuple[CompiledProject, CompiledModel]:
+    """Compile one reshaped-star model over the shared upstream order models."""
+
+    (project_dir / "sqlbuild_project.toml").write_text(
+        _RESHAPED_STAR_PROJECT_TOML, encoding="utf-8"
+    )
+    models: Path = project_dir / "models"
+    models.mkdir()
+    for name, sql in RESHAPED_STAR_UPSTREAM_MODELS.items():
+        (models / f"{name}.sql").write_text(f"MODEL (materialized table);\n{sql}", encoding="utf-8")
+    (models / "reshaped_orders.sql").write_text(
+        f"MODEL (materialized table);\n{query_sql}", encoding="utf-8"
+    )
+    project: CompiledProject = compile_project(
+        discovered_inputs=discover_project_inputs(project_dir=project_dir),
+        adapter=DuckDbAdapter(),
+    )
+    models_by_name: dict[str, CompiledModel] = {model.name: model for model in project.models}
+    return project, models_by_name["reshaped_orders"]
+
+
+def lineage_source_pairs(column: CompiledLineageColumnFact) -> frozenset[tuple[str, str]]:
+    """Return one lineage column's upstream (resource, column) pairs."""
+
+    return frozenset(
+        (source.resource_name, source.column_name) for source in column.upstream_columns
+    )

tests/integration/src/sqlbuild/compiler/pipeline/test_reshaped_star_lineage.py

@@ -0,0 +1,99 @@
+"""Compile star projections over PIVOT and UNPIVOT sources against DuckDB results."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Sequence
+from pathlib import Path
+
+import duckdb
+import pytest
+
+from sqlbuild.compiler.compile.models import CompiledLineageColumnFact
+from tests.integration.src.sqlbuild.compiler.pipeline._test_types import (
+    ReshapedStarIntegrationTestCase,
+    ReshapedStarLineageIntegrationTestCase,
+)
+from tests.integration.src.sqlbuild.compiler.pipeline.helpers import (
+    RESHAPED_STAR_UPSTREAM_MODELS,
+    compile_reshaped_star_model,
+    lineage_source_pairs,
+)
+
+_REF_PATTERN: re.Pattern[str] = re.compile(r'__ref\("(\w+)"\)')
+_PIVOT_OVER_CTE_SQL: str = (
+    "WITH pivot_input AS (SELECT customer_id, category, amount "
+    'FROM __ref("staged_orders")) '
+    "SELECT * FROM pivot_input PIVOT (MAX(amount) FOR category IN ('books', 'games'))"
+)
+_UNPIVOT_OVER_CTE_SQL: str = (
+    'WITH wide AS (SELECT customer_id, books, games FROM __ref("wide_orders")) '
+    "SELECT * FROM wide UNPIVOT (amount FOR category IN (books, games))"
+)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    (
+        ReshapedStarIntegrationTestCase(
+            description="pivot over cte",
+            query_sql=_PIVOT_OVER_CTE_SQL,
+            expected_columns=("customer_id", "books", "games"),
+        ),
+        ReshapedStarIntegrationTestCase(
+            description="unpivot over cte",
+            query_sql=_UNPIVOT_OVER_CTE_SQL,
+            expected_columns=("customer_id", "category", "amount"),
+        ),
+    ),
+    ids=lambda case: case.description,
+)
+def test_given_star_over_reshaped_cte_when_compiling_then_columns_match_duckdb(
+    test_case: ReshapedStarIntegrationTestCase,
+    tmp_path: Path,
+) -> None:
+    project, model = compile_reshaped_star_model(
+        project_dir=tmp_path, query_sql=test_case.query_sql
+    )
+    with duckdb.connect() as connection:
+        for name, sql in RESHAPED_STAR_UPSTREAM_MODELS.items():
+            connection.execute(f"CREATE TABLE {name} AS {sql}")
+        duckdb_columns: tuple[str, ...] = tuple(
+            column[0]
+            for column in connection.execute(
+                _REF_PATTERN.sub(r"\1", test_case.query_sql)
+            ).description
+        )
+
+    assert project.diagnostics == ()
+    assert duckdb_columns == test_case.expected_columns
+    assert model.inferred_columns is not None
+    assert tuple(column.name for column in model.inferred_columns) == test_case.expected_columns
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    (
+        ReshapedStarLineageIntegrationTestCase(
+            description="pivot over cte",
+            query_sql=_PIVOT_OVER_CTE_SQL,
+            expected_sources=(
+                frozenset({("staged_orders", "customer_id")}),
+                frozenset({("staged_orders", "amount")}),
+                frozenset({("staged_orders", "amount")}),
+            ),
+        ),
+    ),
+    ids=lambda case: case.description,
+)
+def test_given_star_over_pivoted_cte_when_compiling_then_lineage_follows_pivot_inputs(
+    test_case: ReshapedStarLineageIntegrationTestCase,
+    tmp_path: Path,
+) -> None:
+    _project, model = compile_reshaped_star_model(
+        project_dir=tmp_path, query_sql=test_case.query_sql
+    )
+    lineage: Sequence[CompiledLineageColumnFact] | None = model.fast_lineage_columns
+
+    assert lineage is not None
+    assert tuple(lineage_source_pairs(column) for column in lineage) == test_case.expected_sources