fix: keep set operations after a parenthesised FROM query at the outer query

Adopt the Polyglot parser fix that stops a parenthesised query in FROM
from absorbing a following UNION, INTERSECT or EXCEPT. Such models no
longer raise false column-count errors, their star columns trace to every
set-operation branch, and the formatter no longer nests the set operation
inside the FROM query.

Add end-to-end compile coverage for DuckDB, Snowflake, BigQuery and
PostgreSQL, plus a DuckDB build case.

Author: kvlonge

Cargo.lock

@@ -64,14 +64,14 @@ junction = { version = "2.1.0", default-features = false }
 mimalloc = { version = "0.1", default-features = false, features = ["local_dynamic_tls"] }
 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 = "ca9e7b6e405fb6f5484472bb08780268789c5e04", default-features = false, features = [
+polyglot-sql = { git = "https://github.com/chio-labs/polyglot.git", rev = "b18dc647e797e480d2faadf3af407077f8e063df", default-features = false, features = [
     "ast-tools",
     "builder",
     "all-dialects",
     "semantic",
     "stacker",
 ] }
-polyglot-sql-function-catalogs = { git = "https://github.com/chio-labs/polyglot.git", rev = "ca9e7b6e405fb6f5484472bb08780268789c5e04", default-features = false, features = [
+polyglot-sql-function-catalogs = { git = "https://github.com/chio-labs/polyglot.git", rev = "b18dc647e797e480d2faadf3af407077f8e063df", default-features = false, features = [
     "dialect-duckdb",
 ] }
 blake3 = { version = "1", features = ["rayon"] }

Cargo.toml

@@ -538,6 +538,17 @@ class SetOperationLifecycleTestCase:
     expected_indented_set_operation_lines: tuple[str, ...]
 
 
+@dataclass(frozen=True)
+class FromQuerySetOperationTestCase:
+    description: str
+    adapter: str
+    models: tuple[SetOperationModel, ...]
+    expected_column_counts: tuple[tuple[str, int], ...]
+    lineage_target: str
+    expected_lineage: tuple[tuple[str, str, str], ...]
+    expected_rule_summary: str
+
+
 @dataclass(frozen=True)
 class SetOperationArityMismatchTestCase:
     description: str

tests/e2e/src/sqlbuild/cli/commands/main/compile/_test_types.py

@@ -2879,15 +2879,21 @@ class SetOperationLifecycleResult(NamedTuple):
 
 
 def write_set_operation_project(
-    *, project_dir: Path, models: tuple[SetOperationModel, ...]
+    *, project_dir: Path, models: tuple[SetOperationModel, ...], adapter: str = "duckdb"
 ) -> Path:
-    """Write a DuckDB project with one view per model and return its database path."""
+    """Write a project with one view per model and return its DuckDB database path.
+
+    Other adapters compile offline against a default database and never connect.
+    """
 
     database: Path = project_dir / "orders.duckdb"
+    target: str = {
+        True: f'[connection]\ndatabase = "{database}"\n',
+        False: '[defaults]\ndatabase = "example"\n',
+    }[adapter == "duckdb"]
     files: dict[str, str] = {
         "sqlbuild_project.toml": (
-            f'name = "orders"\nadapter = "duckdb"\n[connection]\ndatabase = "{database}"\n'
-            '[rules]\nselect = ["SQBRSQL020"]\n'
+            f'name = "orders"\nadapter = "{adapter}"\n{target}[rules]\nselect = ["SQBRSQL020"]\n'
         ),
     }
     for model in models:
@@ -2938,25 +2944,56 @@ def run_set_operation_lifecycle(
                 f"SELECT * FROM analytics.{model.name}"
             )
             relation_shapes[model.name] = (len(relation.columns), len(relation.fetchall()))
-    rules: subprocess.CompletedProcess[str] = run_sqb(
-        project_dir=project_dir, command=("rules", "run", "SQBRSQL020")
+    rules: subprocess.CompletedProcess[str] = run_set_arity_rule(project_dir=project_dir)
+    formatted: subprocess.CompletedProcess[str]
+    indented_set_operation_lines: tuple[str, ...]
+    formatted, indented_set_operation_lines = format_set_operation_models(project_dir=project_dir)
+    return SetOperationLifecycleResult(
+        compiled=compiled,
+        build=build,
+        relation_shapes=relation_shapes,
+        rules=rules,
+        format=formatted,
+        indented_set_operation_lines=indented_set_operation_lines,
+        recompiled=compile_set_operation_project(project_dir=project_dir),
     )
+
+
+def run_set_arity_rule(*, project_dir: Path) -> subprocess.CompletedProcess[str]:
+    """Run the set-operation arity Rule over every model."""
+
+    return run_sqb(project_dir=project_dir, command=("rules", "run", "SQBRSQL020"))
+
+
+def format_set_operation_models(
+    *, project_dir: Path
+) -> tuple[subprocess.CompletedProcess[str], tuple[str, ...]]:
+    """Format every model and return the result and any set-operation lines left indented."""
+
     models_dir: Path = project_dir / "models"
     formatted: subprocess.CompletedProcess[str] = run_sqb(
         project_dir=project_dir, command=("format", str(models_dir))
     )
     indented_set_operation_lines: list[str] = []
     for path in sorted(models_dir.glob("*.sql")):
         indented_set_operation_lines.extend(_indented_set_operation_lines(path))
-    return SetOperationLifecycleResult(
-        compiled=compiled,
-        build=build,
-        relation_shapes=relation_shapes,
-        rules=rules,
-        format=formatted,
-        indented_set_operation_lines=tuple(indented_set_operation_lines),
-        recompiled=compile_set_operation_project(project_dir=project_dir),
+    return formatted, tuple(indented_set_operation_lines)
+
+
+def column_lineage_edges(
+    *, project_dir: Path, target: str
+) -> tuple[subprocess.CompletedProcess[str], tuple[tuple[str, str, str], ...]]:
+    """Return the lineage result and sorted `(model, column, transform)` upstream edges."""
+
+    result: subprocess.CompletedProcess[str] = run_sqb(
+        project_dir=project_dir, command=("lineage", target, "--json")
     )
+    payload: dict[str, Any] = json.loads(result.stdout or '{"trace": []}')
+    edges: list[tuple[str, str, str]] = [
+        (edge["source"]["resource_name"], edge["source"]["column_name"], edge["transform"])
+        for edge in payload["trace"]
+    ]
+    return result, tuple(sorted(edges))
 
 
 class UnionFixtureCompileMeasurement(NamedTuple):

tests/e2e/src/sqlbuild/cli/commands/main/compile/helpers.py

@@ -2,19 +2,24 @@
 
 from __future__ import annotations
 
+import subprocess
 from pathlib import Path
 
 import pytest
 
 from tests.e2e.src.sqlbuild.cli.commands.main.compile._test_types import (
+    FromQuerySetOperationTestCase,
     SetOperationArityMismatchTestCase,
     SetOperationLifecycleTestCase,
     SetOperationModel,
 )
 from tests.e2e.src.sqlbuild.cli.commands.main.compile.helpers import (
     SetOperationCompileResult,
     SetOperationLifecycleResult,
+    column_lineage_edges,
     compile_set_operation_project,
+    format_set_operation_models,
+    run_set_arity_rule,
     run_set_operation_lifecycle,
     write_set_operation_project,
 )
@@ -25,6 +30,43 @@
     "), limits AS (SELECT 10 AS max_amount)\n"
 )
 _ALL_ORDERS: str = "SELECT o.order_id, o.amount FROM orders AS o"
+_ORDERS_QUERY: str = 'SELECT o.order_id, o.amount, o.status FROM __ref("orders") AS o'
+_RETURNS_QUERY: str = 'SELECT r.order_id, r.amount FROM __ref("returns") AS r'
+_FROM_QUERY_MODELS: tuple[SetOperationModel, ...] = (
+    SetOperationModel(
+        description="three-column orders",
+        name="orders",
+        query_sql="SELECT 1 AS order_id, 20 AS amount, 'open' AS status",
+    ),
+    SetOperationModel(
+        description="two-column returns",
+        name="returns",
+        query_sql="SELECT 2 AS order_id, 5 AS amount",
+    ),
+    SetOperationModel(
+        description="narrowed FROM query before UNION ALL",
+        name="union_all_orders",
+        query_sql=f"SELECT order_id, amount FROM ({_ORDERS_QUERY})\nUNION ALL\n{_RETURNS_QUERY}",
+    ),
+    SetOperationModel(
+        description="narrowed FROM query before EXCEPT",
+        name="except_orders",
+        query_sql=f"SELECT order_id, amount FROM ({_ORDERS_QUERY})\nEXCEPT\n{_RETURNS_QUERY}",
+    ),
+    SetOperationModel(
+        description="narrowed FROM query before INTERSECT",
+        name="intersect_orders",
+        query_sql=f"SELECT order_id, amount FROM ({_ORDERS_QUERY})\nINTERSECT\n{_RETURNS_QUERY}",
+    ),
+    SetOperationModel(
+        description="star over a FROM query before UNION ALL",
+        name="starred_orders",
+        query_sql=(
+            'SELECT * FROM (SELECT o.order_id, o.amount FROM __ref("orders") AS o)\n'
+            f"UNION ALL\n{_RETURNS_QUERY}"
+        ),
+    ),
+)
 
 
 @pytest.mark.parametrize(
@@ -128,6 +170,16 @@
                     ),
                     expected_row_count=4,
                 ),
+                SetOperationModel(
+                    description="narrowed FROM query before UNION ALL",
+                    name="from_query_union_all",
+                    query_sql=(
+                        f"{_INPUTS}SELECT order_id, amount\n"
+                        "FROM (SELECT o.order_id, o.amount, 'open' AS status FROM orders AS o)\n"
+                        f"UNION ALL\n{_ALL_ORDERS}"
+                    ),
+                    expected_row_count=4,
+                ),
                 SetOperationModel(
                     description="parenthesized branches around a scalar subquery",
                     name="parenthesized_branches_union_all",
@@ -141,7 +193,7 @@
             ),
             expected_column_count=2,
             expected_diagnostics=(),
-            expected_rule_summary="10 models evaluated, 0 findings",
+            expected_rule_summary="11 models evaluated, 0 findings",
             expected_indented_set_operation_lines=(),
         )
     ],
@@ -175,6 +227,63 @@ def test_given_set_operation_after_subquery_operand_when_compiling_building_and_
     assert result.recompiled.column_counts == expected_column_counts
 
 
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        FromQuerySetOperationTestCase(
+            description=f"{adapter} FROM query before a set operation",
+            adapter=adapter,
+            models=_FROM_QUERY_MODELS,
+            expected_column_counts=(
+                ("union_all_orders", 2),
+                ("except_orders", 2),
+                ("intersect_orders", 2),
+            ),
+            lineage_target="starred_orders.order_id",
+            expected_lineage=(("orders", "order_id", "direct"), ("returns", "order_id", "direct")),
+            expected_rule_summary="6 models evaluated, 0 findings",
+        )
+        for adapter in ("duckdb", "snowflake", "bigquery", "postgres")
+    ],
+    ids=lambda case: case.description,
+)
+def test_given_set_operation_after_parenthesized_from_query_when_compiling_then_outer_queries_combine(
+    test_case: FromQuerySetOperationTestCase,
+    tmp_path: Path,
+) -> None:
+    project_dir: Path = tmp_path / "orders_project"
+    write_set_operation_project(
+        project_dir=project_dir, models=test_case.models, adapter=test_case.adapter
+    )
+
+    compiled: SetOperationCompileResult = compile_set_operation_project(project_dir=project_dir)
+    lineage: subprocess.CompletedProcess[str]
+    lineage_edges: tuple[tuple[str, str, str], ...]
+    lineage, lineage_edges = column_lineage_edges(
+        project_dir=project_dir, target=test_case.lineage_target
+    )
+    rules: subprocess.CompletedProcess[str] = run_set_arity_rule(project_dir=project_dir)
+    formatted: subprocess.CompletedProcess[str]
+    indented_set_operation_lines: tuple[str, ...]
+    formatted, indented_set_operation_lines = format_set_operation_models(project_dir=project_dir)
+    recompiled: SetOperationCompileResult = compile_set_operation_project(project_dir=project_dir)
+
+    expected_column_counts: dict[str, int] = dict(test_case.expected_column_counts)
+    assert compiled.exit_code == 0, compiled.output
+    assert compiled.diagnostics == ()
+    assert {
+        name: compiled.column_counts[name] for name in expected_column_counts
+    } == expected_column_counts
+    assert lineage.returncode == 0, lineage.stdout + lineage.stderr
+    assert lineage_edges == test_case.expected_lineage
+    assert rules.returncode == 0, rules.stdout + rules.stderr
+    assert test_case.expected_rule_summary in rules.stdout + rules.stderr
+    assert formatted.returncode == 0, formatted.stdout + formatted.stderr
+    assert indented_set_operation_lines == ()
+    assert recompiled.exit_code == 0, recompiled.output
+    assert recompiled.diagnostics == ()
+
+
 @pytest.mark.parametrize(
     "test_case",
     [