fix(test): ignore comments around set-operation branches in SQL-test fixtures

A comment before or after UNION ALL in an __expected__ fixture stayed
attached to the neighbouring branch, so the projection scanners read the
comment as an alias and reported a false P001 branch mismatch. Branch and
projection splits now drop surrounding comments, alias detection ignores
trailing comments, the Python fallback splitter follows the adapter's
lexical syntax, and Polyglot's annotated comment wrapper is unwrapped.

Author: kvlonge

crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs

@@ -512,7 +512,7 @@ fn projection_name(expression: &str, file: &str, label: &str) -> Result<String,
     if let Some(position) = find_last_top_level_keyword(expression, "AS")? {
         let alias_start = skip_ignorable(expression, position + 2)?;
         if let Some((alias, end)) = read_identifier(expression, alias_start)
-            && expression[end..].trim().is_empty()
+            && skip_ignorable(expression, end)? == expression.len()
         {
             return Ok(alias);
         }
@@ -880,8 +880,34 @@ fn scan_code<T>(
     Ok(None)
 }
 
-fn non_empty_trimmed(value: &str) -> Option<&str> {
-    Some(value.trim()).filter(|trimmed| !trimmed.is_empty())
+/// Return `value` without surrounding whitespace and comments, or `None` when no code remains.
+fn non_empty_trimmed(value: &str) -> Result<Option<&str>, String> {
+    if !value.contains("--") && !value.contains("/*") {
+        return Ok(Some(value.trim()).filter(|trimmed| !trimmed.is_empty()));
+    }
+    let start = skip_ignorable(value, 0)?;
+    let mut end = start;
+    let mut index = start;
+    while index < value.len() {
+        if let Some(comment) = comment_end(value.as_bytes(), index)
+            .map_err(|error| scan_error_message(error, "SQL test"))?
+        {
+            index = comment;
+            continue;
+        }
+        let next = skip_non_code(value, index)?;
+        if next != index {
+            end = next;
+            index = next;
+            continue;
+        }
+        let width = char_len(value, index);
+        if !value[index..].starts_with(char::is_whitespace) {
+            end = index + width;
+        }
+        index += width;
+    }
+    Ok(Some(&value[start..end]).filter(|trimmed| !trimmed.is_empty()))
 }
 
 /// Split on top-level `UNION`, `INTERSECT` and `EXCEPT`, each with an optional `ALL`/`DISTINCT`.
@@ -895,7 +921,7 @@ pub(crate) fn split_set_operations(sql: &str) -> Result<Vec<&str>, String> {
         let Some(end) = set_operator_end(sql, index)? else {
             return Ok(CodeStep::Advance);
         };
-        values.extend(non_empty_trimmed(&sql[start..index]));
+        values.extend(non_empty_trimmed(&sql[start..index])?);
         let mut resume = skip_ignorable(sql, end)?;
         if let Some(quantifier_end) =
             consume_keyword(sql, resume, "ALL").or_else(|| consume_keyword(sql, resume, "DISTINCT"))
@@ -905,7 +931,7 @@ pub(crate) fn split_set_operations(sql: &str) -> Result<Vec<&str>, String> {
         start = resume;
         Ok(CodeStep::Resume(resume))
     })?;
-    values.extend(non_empty_trimmed(&sql[start..]));
+    values.extend(non_empty_trimmed(&sql[start..])?);
     Ok(values)
 }
 
@@ -927,12 +953,12 @@ pub(crate) fn split_top_level(sql: &str, separator: u8) -> Result<Vec<&str>, Str
     let mut start = 0;
     scan_code::<()>(sql, 0, |index, depth| {
         if depth == 0 && byte_at(sql, index) == Some(separator) {
-            values.extend(non_empty_trimmed(&sql[start..index]));
+            values.extend(non_empty_trimmed(&sql[start..index])?);
             start = index + 1;
         }
         Ok(CodeStep::Advance)
     })?;
-    values.extend(non_empty_trimmed(&sql[start..]));
+    values.extend(non_empty_trimmed(&sql[start..])?);
     Ok(values)
 }
 

crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_extraction.rs

@@ -91,12 +91,26 @@ fn given_expected_cte_sql_when_scanning_top_level_then_set_operations_commas_and
             ]),
             expected_commas: Ok(vec![
                 "SELECT 'x, UNION (' AS a",
-                "/* UNION , FROM */ \"b)\"",
+                "\"b)\"",
                 "`c,d`",
                 "f(1, (SELECT 2 FROM t UNION SELECT 3)) -- , FROM\nFROM t",
             ]),
             expected_from: Ok(Some(107)),
         },
+        TopLevelScanTestCase {
+            description: "comments around branches and projections are not code",
+            sql: "-- lead\nSELECT 1 AS a, /* b */ 2 AS b -- first\n/* x */ UNION ALL -- second\nSELECT '--' AS a, 3 AS b /* end */",
+            expected_unions: Ok(vec![
+                "SELECT 1 AS a, /* b */ 2 AS b",
+                "SELECT '--' AS a, 3 AS b",
+            ]),
+            expected_commas: Ok(vec![
+                "SELECT 1 AS a",
+                "2 AS b -- first\n/* x */ UNION ALL -- second\nSELECT '--' AS a",
+                "3 AS b",
+            ]),
+            expected_from: Ok(None),
+        },
         TopLevelScanTestCase {
             description: "keyword boundaries and empty pieces",
             sql: ",SELECT reunion, union_id,, x FROMAGE UNION",

src/sqlbuild/compiler/compile/_helpers/analysis/tests.py

@@ -15,14 +15,19 @@
     split_set_operation_branches,
 )
 from sqlbuild.compiler.sql_analysis.main.import_polyglot_sql import import_polyglot_sql
+from sqlbuild.compiler.sql_analysis.models import SqlLexicalSyntax
 
 _POLYGLOT_VALUES_EXPRESSION_NAME: str = "Values"
 _POLYGLOT_VALUES_SET_ALIAS: str = "_values"
 _SCAN_CONTEXT: str = "SQL test"
 
 
 def extract_expected_branch_column_names_with_sql_analysis(
-    *, sql: str, file_label: str, label: str = "__expected__<model>"
+    *,
+    sql: str,
+    file_label: str,
+    syntax: SqlLexicalSyntax,
+    label: str = "__expected__<model>",
 ) -> tuple[tuple[str, ...], ...] | None:
     """Return expected SELECT branch names using required Polyglot analysis."""
 
@@ -32,17 +37,23 @@ def extract_expected_branch_column_names_with_sql_analysis(
         parsed_expression: Any = polyglot_module.parse_one(sql, dialect="generic")
     except polyglot_module.PolyglotError:
         return _extract_unparsed_branch_names(
-            sql=sql, polyglot_module=polyglot_module, file_label=file_label, label=label
+            sql=sql,
+            polyglot_module=polyglot_module,
+            file_label=file_label,
+            label=label,
+            syntax=syntax,
         )
     return _extract_branch_names(expression=parsed_expression, file_label=file_label, label=label)
 
 
 def _extract_unparsed_branch_names(
-    *, sql: str, polyglot_module: Any, file_label: str, label: str
+    *, sql: str, polyglot_module: Any, file_label: str, label: str, syntax: SqlLexicalSyntax
 ) -> tuple[tuple[str, ...], ...] | None:
     """Parse each top-level set-operation branch separately when the whole query cannot parse."""
 
-    branches: tuple[str, ...] = split_set_operation_branches(sql=sql, context=_SCAN_CONTEXT)
+    branches: tuple[str, ...] = split_set_operation_branches(
+        sql=sql, context=_SCAN_CONTEXT, syntax=syntax
+    )
     if len(branches) <= 1:
         return None
     try:

src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py

@@ -762,12 +762,14 @@ def _extract_expected_branch_column_names(
 ) -> tuple[tuple[str, ...], ...]:
     sql_analysis_column_names: tuple[tuple[str, ...], ...] | None = (
         extract_expected_branch_column_names_with_sql_analysis(
-            sql=sql, file_label=file_label, label=label
+            sql=sql, file_label=file_label, label=label, syntax=syntax
         )
     )
     if sql_analysis_column_names is not None:
         return sql_analysis_column_names
-    branches: tuple[str, ...] = split_set_operation_branches(sql=sql, context=_CONTEXT)
+    branches: tuple[str, ...] = split_set_operation_branches(
+        sql=sql, context=_CONTEXT, syntax=syntax
+    )
     return tuple(
         _extract_expected_select_column_names(
             branch_sql=branch, file_label=file_label, label=label, syntax=syntax
@@ -837,9 +839,13 @@ def _extract_expected_projection_name(
     alias_name: str | None = _extract_as_alias(expression=expression, syntax=syntax)
     if alias_name is not None:
         return alias_name
-    stripped_expression: str = expression.strip()
-    if _is_simple_identifier(stripped_expression):
-        return stripped_expression
+    start: int = _skip_ignorable(sql=expression, start=0, syntax=syntax)
+    if start < len(expression) and is_identifier_start(expression[start]):
+        name: str
+        end: int
+        name, end = _read_identifier(sql=expression, start=start, file_label="projection")
+        if _skip_ignorable(sql=expression, start=end, syntax=syntax) == len(expression):
+            return name
     raise CompileInputError(
         f"SQL test '{file_label}' must alias every non-trivial {label} projection"
     )
@@ -862,17 +868,11 @@ def _extract_as_alias(*, expression: str, syntax: SqlLexicalSyntax) -> str | Non
                 start=alias_index,
                 file_label="projection",
             )
-            if not expression[alias_end:].strip():
+            if _skip_ignorable(sql=expression, start=alias_end, syntax=syntax) == len(expression):
                 last_alias_name = alias_name
     return last_alias_name
 
 
-def _is_simple_identifier(value: str) -> bool:
-    if not value or not is_identifier_start(value[0]):
-        return False
-    return all(is_identifier_character(character) for character in value[1:])
-
-
 def _contains_select_star(*, sql: str, syntax: SqlLexicalSyntax) -> bool:
     index: int
     for index, _depth in iter_code_positions(sql=sql, context=_CONTEXT, syntax=syntax):

src/sqlbuild/compiler/compile/constants.py

@@ -96,7 +96,7 @@
 )
 POLYGLOT_SELECT_EXPRESSION_NAME: str = "Select"
 POLYGLOT_COLUMN_EXPRESSION_NAME: str = "Column"
-POLYGLOT_WRAPPER_EXPRESSION_NAMES: frozenset[str] = frozenset({"Subquery", "Paren"})
+POLYGLOT_WRAPPER_EXPRESSION_NAMES: frozenset[str] = frozenset({"Annotated", "Subquery", "Paren"})
 
 TABLE_FUNCTION_RETURN_KEYS: frozenset[str] = frozenset({"table"})
 

src/sqlbuild/compiler/sql_analysis/_helpers/scanning.py

@@ -231,6 +231,50 @@ def iter_code_positions_impl(
         index += 1
 
 
+def skip_ignorable_impl(
+    *, sql: str, start: int, context: str = "SQL", syntax: SqlLexicalSyntax | None = None
+) -> int:
+    """Return the first offset at or after `start` that is not whitespace or a comment."""
+
+    index: int = start
+    while index < len(sql):
+        if sql[index].isspace():
+            index += 1
+            continue
+        if sql[index] in SQL_TEXT_START_CHARACTERS:
+            return index
+        comment_end: int | None = _non_code_end(
+            sql=sql, index=index, context=context, syntax=syntax
+        )
+        if comment_end is None:
+            return index
+        index = comment_end
+    return index
+
+
+def strip_ignorable_impl(
+    *, sql: str, context: str = "SQL", syntax: SqlLexicalSyntax | None = None
+) -> str:
+    """Return the SQL without leading or trailing whitespace and comments."""
+
+    code_start: int = skip_ignorable_impl(sql=sql, start=0, context=context, syntax=syntax)
+    code_end: int = code_start
+    index: int = code_start
+    while index < len(sql):
+        non_code_end: int | None = _non_code_end(
+            sql=sql, index=index, context=context, syntax=syntax
+        )
+        if non_code_end is None:
+            index += 1
+            if not sql[index - 1].isspace():
+                code_end = index
+            continue
+        if sql[index] in SQL_TEXT_START_CHARACTERS:
+            code_end = non_code_end
+        index = non_code_end
+    return sql[code_start:code_end]
+
+
 def _non_code_end(
     *, sql: str, index: int, context: str, syntax: SqlLexicalSyntax | None
 ) -> int | None:

src/sqlbuild/compiler/sql_analysis/_helpers/set_operations.py

@@ -5,9 +5,10 @@
 from sqlbuild.compiler.sql_analysis._helpers.scanning import (
     is_identifier_character_impl,
     iter_code_positions_impl,
-    skip_block_comment_impl,
-    skip_line_comment_impl,
+    skip_ignorable_impl,
+    strip_ignorable_impl,
 )
+from sqlbuild.compiler.sql_analysis.models import SqlLexicalSyntax
 
 _SET_OPERATOR_KEYWORDS: tuple[str, ...] = ("UNION", "INTERSECT")
 _EXCEPT_KEYWORD: str = "EXCEPT"
@@ -16,8 +17,10 @@
 _STAR_CHARACTER: str = "*"
 
 
-def split_set_operation_branches_impl(*, sql: str, context: str) -> tuple[str, ...]:
-    """Split SQL on top-level `UNION`, `INTERSECT` and `EXCEPT` with optional quantifiers."""
+def split_set_operation_branches_impl(
+    *, sql: str, context: str, syntax: SqlLexicalSyntax | None = None
+) -> tuple[str, ...]:
+    """Split SQL on top-level set operators; branches exclude surrounding comments."""
 
     branches: list[str] = []
     branch_start: int = 0
@@ -26,7 +29,7 @@ def split_set_operation_branches_impl(*, sql: str, context: str) -> tuple[str, .
     previous_depth: int = 0
     index: int
     depth: int
-    for index, depth in iter_code_positions_impl(sql=sql, context=context):
+    for index, depth in iter_code_positions_impl(sql=sql, context=context, syntax=syntax):
         if depth != previous_depth:
             previous_code = None
         previous_depth = depth
@@ -39,17 +42,23 @@ def split_set_operation_branches_impl(*, sql: str, context: str) -> tuple[str, .
             previous_code = sql[index]
         if operator_end is None:
             continue
-        branch_sql: str = sql[branch_start:index].strip()
+        branch_sql: str = strip_ignorable_impl(
+            sql=sql[branch_start:index], context=context, syntax=syntax
+        )
         if branch_sql:
             branches.append(branch_sql)
-        resume = _skip_ignorable(sql=sql, start=operator_end, context=context)
+        resume = skip_ignorable_impl(sql=sql, start=operator_end, context=context, syntax=syntax)
         quantifier_end: int | None = _keyword_end(
             sql=sql, start=resume, keyword=_ALL_KEYWORD
         ) or _keyword_end(sql=sql, start=resume, keyword=_DISTINCT_KEYWORD)
         if quantifier_end is not None:
-            resume = _skip_ignorable(sql=sql, start=quantifier_end, context=context)
+            resume = skip_ignorable_impl(
+                sql=sql, start=quantifier_end, context=context, syntax=syntax
+            )
         branch_start = resume
-    final_branch_sql: str = sql[branch_start:].strip()
+    final_branch_sql: str = strip_ignorable_impl(
+        sql=sql[branch_start:], context=context, syntax=syntax
+    )
     if final_branch_sql:
         branches.append(final_branch_sql)
     return tuple(branches)
@@ -77,19 +86,3 @@ def _keyword_end(*, sql: str, start: int, keyword: str) -> int | None:
     if start > 0 and is_identifier_character_impl(sql[start - 1]):
         return None
     return keyword_end
-
-
-def _skip_ignorable(*, sql: str, start: int, context: str) -> int:
-    index: int = start
-    while index < len(sql):
-        if sql[index].isspace():
-            index += 1
-            continue
-        if sql.startswith("--", index):
-            index = skip_line_comment_impl(sql=sql, start=index)
-            continue
-        if sql.startswith("/*", index):
-            index = skip_block_comment_impl(sql=sql, start=index, context=context)
-            continue
-        return index
-    return index

src/sqlbuild/compiler/sql_analysis/main/_split_set_operation_branches.py

@@ -3,9 +3,12 @@
 from sqlbuild.compiler.sql_analysis._helpers.set_operations import (
     split_set_operation_branches_impl,
 )
+from sqlbuild.compiler.sql_analysis.models import SqlLexicalSyntax
 
 
-def split_set_operation_branches(*, sql: str, context: str = "SQL") -> tuple[str, ...]:
+def split_set_operation_branches(
+    *, sql: str, context: str = "SQL", syntax: SqlLexicalSyntax | None = None
+) -> tuple[str, ...]:
     """Split SQL on top-level `UNION`, `INTERSECT` and `EXCEPT`, ignoring quotes and comments."""
 
-    return split_set_operation_branches_impl(sql=sql, context=context)
+    return split_set_operation_branches_impl(sql=sql, context=context, syntax=syntax)

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

@@ -212,3 +212,13 @@ class AuthoredSqlParityTestCase:
     description: str
     compiled_test_name: str
     expected_fragments: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class CommentedFixtureE2ETestCase:
+    """SQL tests whose fixture CTEs carry comments around set-operation branches."""
+
+    description: str
+    test_files: dict[str, str]
+    expected_exit_code: int
+    expected_output_fragments: tuple[str, ...]