fix: keep set operations after subquery operands at the outer query

Cargo.toml

@@ -17,7 +17,7 @@ globset = "0.4"
 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 = "25ad011c49f2c7cd48f4427abb8fc0a6af2e4120", default-features = false, features = [
+polyglot-sql = { git = "https://github.com/chio-labs/polyglot.git", rev = "a7413439c520d7949924149b7ea02fa3817463fc", default-features = false, features = [
     "ast-tools",
     "builder",
     "dialect-bigquery",
@@ -29,7 +29,7 @@ polyglot-sql = { git = "https://github.com/chio-labs/polyglot.git", rev = "25ad0
     "semantic",
     "stacker",
 ] }
-polyglot-sql-function-catalogs = { git = "https://github.com/chio-labs/polyglot.git", rev = "25ad011c49f2c7cd48f4427abb8fc0a6af2e4120", default-features = false, features = [
+polyglot-sql-function-catalogs = { git = "https://github.com/chio-labs/polyglot.git", rev = "a7413439c520d7949924149b7ea02fa3817463fc", default-features = false, features = [
     "dialect-duckdb",
 ] }
 quote = "1"

crates/sqlbuild-rules-native/src/sql_lint/tests.rs

@@ -16,6 +16,8 @@ mod helpers;
 mod layout_rules;
 #[path = "tests/test_line_wrap.rs"]
 mod line_wrap;
+#[path = "tests/test_set_operation_formatter.rs"]
+mod set_operation_formatter;
 #[path = "tests/test_sql_lint.rs"]
 mod sql_lint;
 #[path = "tests/test_types.rs"]

crates/sqlbuild-rules-native/src/sql_lint/tests/test_set_operation_formatter.rs

@@ -0,0 +1,59 @@
+use serde_json::{Value, json};
+
+use crate::sql_lint::main::formatter::format_json;
+use crate::sql_lint::tests::test_types;
+
+const DIALECTS: [&str; 6] = [
+    "bigquery",
+    "databricks",
+    "duckdb",
+    "postgres",
+    "snowflake",
+    "tsql",
+];
+
+#[test]
+fn given_set_operation_after_scalar_subquery_operand_when_formatting_then_keyword_stays_outer()
+-> Result<(), String> {
+    let test_cases = [
+        test_types::FormatTestCase {
+            description: "UNION ALL after WHERE scalar subquery",
+            sql: "SELECT o.id, o.total FROM orders AS o WHERE o.total < (SELECT MAX(l.amount) FROM limits AS l) UNION ALL SELECT o.id, o.total FROM orders AS o",
+            expected_sql: "SELECT\n  o.id,\n  o.total\nFROM orders AS o\nWHERE\n  o.total < (\n    SELECT MAX(l.amount)\n    FROM limits AS l\n  )\nUNION ALL\nSELECT\n  o.id,\n  o.total\nFROM orders AS o",
+            expected_changed: true,
+        },
+        test_types::FormatTestCase {
+            description: "EXCEPT after WHERE scalar subquery",
+            sql: "SELECT o.id FROM orders AS o WHERE o.total > (SELECT MAX(l.amount) FROM limits AS l) EXCEPT SELECT o.id FROM orders AS o",
+            expected_sql: "SELECT o.id\nFROM orders AS o\nWHERE\n  o.total > (\n    SELECT MAX(l.amount)\n    FROM limits AS l\n  )\nEXCEPT\nSELECT o.id\nFROM orders AS o",
+            expected_changed: true,
+        },
+        test_types::FormatTestCase {
+            description: "INTERSECT after WHERE scalar subquery",
+            sql: "SELECT o.id FROM orders AS o WHERE o.total > (SELECT MAX(l.amount) FROM limits AS l) INTERSECT SELECT o.id FROM orders AS o",
+            expected_sql: "SELECT o.id\nFROM orders AS o\nWHERE\n  o.total > (\n    SELECT MAX(l.amount)\n    FROM limits AS l\n  )\nINTERSECT\nSELECT o.id\nFROM orders AS o",
+            expected_changed: true,
+        },
+    ];
+
+    for dialect in DIALECTS {
+        for test_case in &test_cases {
+            let response = format_json(
+                &json!({"version": 1, "sql": test_case.sql, "dialect": dialect}).to_string(),
+            )?;
+            let payload: Value =
+                serde_json::from_str(&response).map_err(|error| error.to_string())?;
+            assert_eq!(
+                payload["sql"], test_case.expected_sql,
+                "{dialect}: {}",
+                test_case.description
+            );
+            assert_eq!(
+                payload["changed"], test_case.expected_changed,
+                "{dialect}: {}",
+                test_case.description
+            );
+        }
+    }
+    Ok(())
+}

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

@@ -445,3 +445,41 @@ class RequireSqlAnalysisCase:
     expected_returncode: int
     expected_diagnostics: tuple[tuple[str, str | None, int | None], ...]
     expected_text_fragments: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class SetOperationModel:
+    description: str
+    name: str
+    query_sql: str
+    expected_row_count: int | None = None
+
+
+@dataclass(frozen=True)
+class SetOperationLifecycleTestCase:
+    description: str
+    models: tuple[SetOperationModel, ...]
+    expected_column_count: int
+    expected_diagnostics: tuple[tuple[str, str], ...]
+    expected_rule_summary: str
+    expected_indented_set_operation_lines: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class SetOperationArityMismatchTestCase:
+    description: str
+    models: tuple[SetOperationModel, ...]
+    expected_exit_code: int
+    expected_diagnostics: tuple[tuple[str, str], ...]
+
+
+@dataclass(frozen=True)
+class UnionFixtureCompileScalingTestCase:
+    description: str
+    sql_test_count: int
+    small_fixture_rows: int
+    large_fixture_rows: int
+    measured_runs: int
+    expected_sql_tests: int
+    expected_errors: int
+    expected_max_scaling_ratio: float

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

@@ -19,6 +19,7 @@
 from types import FrameType
 from typing import Any, NamedTuple, cast
 
+import duckdb
 import pytest
 
 from scripts.cold_compile_performance.main.read_compile_measurement import read_compile_measurement
@@ -29,8 +30,9 @@
 from tests.e2e.src.sqlbuild.cli.commands.main.compile._test_types import (
     FreshProcessCompileCachePerformanceGuardTestCase,
     SemanticCorpusCase,
+    SetOperationModel,
 )
-from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import prepare_inline_project
+from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import prepare_inline_project, run_sqb
 
 _DBT_SHAPED_SQL_SIZE_PROFILE: tuple[tuple[float, int], ...] = (
     (0.50, 1_800),
@@ -2601,3 +2603,178 @@ def require_sql_analysis_output(
         tuple((item["code"], item.get("path"), item.get("line")) for item in diagnostics),
         text,
     )
+
+
+_INDENTED_SET_OPERATION_PATTERN: re.Pattern[str] = re.compile(
+    r"^[ \t]+(?:UNION|EXCEPT|INTERSECT)\b.*$", re.MULTILINE
+)
+
+
+class SetOperationCompileResult(NamedTuple):
+    exit_code: int
+    diagnostics: tuple[tuple[str, str], ...]
+    column_counts: dict[str, int]
+    output: str
+
+
+class SetOperationLifecycleResult(NamedTuple):
+    compiled: SetOperationCompileResult
+    build: subprocess.CompletedProcess[str]
+    relation_shapes: dict[str, tuple[int, int]]
+    rules: subprocess.CompletedProcess[str]
+    format: subprocess.CompletedProcess[str]
+    indented_set_operation_lines: tuple[str, ...]
+    recompiled: SetOperationCompileResult
+
+
+def write_set_operation_project(
+    *, project_dir: Path, models: tuple[SetOperationModel, ...]
+) -> Path:
+    """Write a DuckDB project with one view per model and return its database path."""
+
+    database: Path = project_dir / "orders.duckdb"
+    files: dict[str, str] = {
+        "sqlbuild_project.toml": (
+            f'name = "orders"\nadapter = "duckdb"\n[connection]\ndatabase = "{database}"\n'
+            '[rules]\nselect = ["SQBRSQL020"]\n'
+        ),
+    }
+    for model in models:
+        files[f"models/{model.name}.sql"] = (
+            f"MODEL (materialized view, schema analytics);\n\n{model.query_sql}\n"
+        )
+    prepare_inline_project(
+        tmp_path=project_dir.parent, project_name=project_dir.name, repo_files=files
+    )
+    return database
+
+
+def compile_set_operation_project(*, project_dir: Path) -> SetOperationCompileResult:
+    """Compile without cache and return diagnostics and per-model output column counts."""
+
+    result: subprocess.CompletedProcess[str] = run_sqb(
+        project_dir=project_dir, command=("compile", "--no-cache", "--json")
+    )
+    payload: dict[str, Any] = json.loads(result.stdout)
+    return SetOperationCompileResult(
+        exit_code=result.returncode,
+        diagnostics=tuple((item["code"], item["message"]) for item in payload["diagnostics"]),
+        column_counts={
+            model["name"]: model["column_count"] for model in payload["resources"]["models"]
+        },
+        output=result.stdout + result.stderr,
+    )
+
+
+def _indented_set_operation_lines(path: Path) -> list[str]:
+    lines: list[str] = _INDENTED_SET_OPERATION_PATTERN.findall(path.read_text(encoding="utf-8"))
+    return [f"{path.name}: {line}" for line in lines]
+
+
+def run_set_operation_lifecycle(
+    *, project_dir: Path, models: tuple[SetOperationModel, ...]
+) -> SetOperationLifecycleResult:
+    """Compile, build, run Rules, format, and recompile a set-operation project."""
+
+    database: Path = write_set_operation_project(project_dir=project_dir, models=models)
+    compiled: SetOperationCompileResult = compile_set_operation_project(project_dir=project_dir)
+    build: subprocess.CompletedProcess[str] = run_sqb(project_dir=project_dir, command=("build",))
+    relation_shapes: dict[str, tuple[int, int]] = {}
+    with duckdb.connect(str(database), read_only=True) as connection:
+        for model in models:
+            relation: duckdb.DuckDBPyRelation = connection.sql(
+                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")
+    )
+    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),
+    )
+
+
+class UnionFixtureCompileMeasurement(NamedTuple):
+    analysis_native_ms: int
+    sql_tests: int
+    errors: int
+
+
+def _union_all_fixture_rows(*, row_count: int, status: str) -> str:
+    return " UNION ALL\n".join(
+        f"  SELECT {row} AS order_id, CAST({row} AS DOUBLE) AS amount, '{status}' AS status"
+        for row in range(row_count)
+    )
+
+
+def write_union_fixture_test_project(
+    *, project_dir: Path, sql_test_count: int, fixture_row_count: int
+) -> None:
+    """Write a DuckDB project whose SQL tests use long UNION ALL fixture chains."""
+
+    files: dict[str, str] = {
+        "sqlbuild_project.toml": (
+            f'name = "orders"\nadapter = "duckdb"\n[connection]\n'
+            f'database = "{project_dir / "orders.duckdb"}"\n[rules]\nselect = []\n'
+        ),
+        "models/base_orders.sql": (
+            "MODEL (materialized view);\n\n"
+            "SELECT 1 AS order_id, CAST(1 AS DOUBLE) AS amount, 'open' AS status\n"
+        ),
+        "models/orders.sql": (
+            "MODEL (materialized view);\n\n"
+            'SELECT b.order_id, b.amount, b.status FROM __ref("base_orders") AS b\n'
+        ),
+    }
+    for index in range(sql_test_count):
+        fixture_rows: str = _union_all_fixture_rows(
+            row_count=fixture_row_count, status=f"status_{index:02d}"
+        )
+        files[f"tests/unit/orders_case_{index:02d}.sql"] = (
+            f'TEST (name "orders_case_{index:02d}");\n\n'
+            f"WITH\n__ref__base_orders AS (\n{fixture_rows}\n),\n"
+            f"__expected__orders AS (\n{fixture_rows}\n)\nSELECT 1\n"
+        )
+    prepare_inline_project(
+        tmp_path=project_dir.parent, project_name=project_dir.name, repo_files=files
+    )
+
+
+def measure_union_fixture_compile(
+    *, project_dir: Path, runs: int
+) -> UnionFixtureCompileMeasurement:
+    """Return the fastest native analysis time of fresh uncached compiles."""
+
+    skip_actions: dict[bool, Callable[[], None]] = {
+        False: _continue_compile_benchmark,
+        True: _skip_compile_benchmark,
+    }
+    skip_actions[os.environ.get("SQLBUILD_SKIP_PERFORMANCE_TESTS") == "1"]()
+    measurements: list[UnionFixtureCompileMeasurement] = []
+    for _ in range(runs):
+        result: subprocess.CompletedProcess[str] = run_installed_sqb(
+            project_dir=project_dir, args=("compile", "--no-cache", "--json"), env={}
+        )
+        assert result.returncode == 0, result.stdout + result.stderr
+        payload: dict[str, Any] = json.loads(result.stdout)
+        measurements.append(
+            UnionFixtureCompileMeasurement(
+                analysis_native_ms=payload["compile_timings"]["analysis_native_ms"],
+                sql_tests=payload["summary"]["tests"],
+                errors=payload["summary"]["errors"],
+            )
+        )
+    return min(measurements)