Add comment preservation to parser for format roundtrip tests

ast/ast.go

@@ -14,6 +14,63 @@ type Node interface {
 	End() token.Position
 }
 
+// Comment represents a SQL comment with its position.
+type Comment struct {
+	Position token.Position `json:"-"`
+	Text     string         `json:"text"`
+}
+
+func (c *Comment) Pos() token.Position { return c.Position }
+func (c *Comment) End() token.Position { return c.Position }
+
+// EndsOnLine returns true if the comment ends on or before the given line.
+func (c *Comment) EndsOnLine(line int) bool {
+	endLine := c.Position.Line
+	for _, r := range c.Text {
+		if r == '\n' {
+			endLine++
+		}
+	}
+	return endLine <= line
+}
+
+// StatementWithComments wraps a statement with its associated comments.
+type StatementWithComments struct {
+	Statement        Statement  `json:"-"`
+	LeadingComments  []*Comment `json:"-"`
+	TrailingComments []*Comment `json:"-"`
+}
+
+// MarshalJSON delegates JSON serialization to the wrapped statement.
+func (s *StatementWithComments) MarshalJSON() ([]byte, error) {
+	return json.Marshal(s.Statement)
+}
+
+func (s *StatementWithComments) Pos() token.Position {
+	if s.Statement != nil {
+		return s.Statement.Pos()
+	}
+	return token.Position{}
+}
+
+func (s *StatementWithComments) End() token.Position {
+	if s.Statement != nil {
+		return s.Statement.End()
+	}
+	return token.Position{}
+}
+
+func (s *StatementWithComments) statementNode() {}
+
+// UnwrapStatement extracts the underlying statement from a StatementWithComments,
+// or returns the statement as-is if it's not wrapped.
+func UnwrapStatement(stmt Statement) Statement {
+	if wrapped, ok := stmt.(*StatementWithComments); ok {
+		return wrapped.Statement
+	}
+	return stmt
+}
+
 // Statement is the interface implemented by all statement nodes.
 type Statement interface {
 	Node

internal/explain/explain.go

@@ -32,6 +32,11 @@ func Node(sb *strings.Builder, node interface{}, depth int) {
 	indent := strings.Repeat(" ", depth)
 
 	switch n := node.(type) {
+	// Handle statement with comments wrapper - unwrap and explain the inner statement
+	case *ast.StatementWithComments:
+		Node(sb, n.Statement, depth)
+		return
+
 	// Select statements
 	case *ast.SelectWithUnionQuery:
 		explainSelectWithUnionQuery(sb, n, indent, depth)

internal/format/format.go

@@ -20,12 +20,36 @@ func Format(stmts []ast.Statement) string {
 	return sb.String()
 }
 
+// formatComments writes comments to the builder.
+func formatComments(sb *strings.Builder, comments []*ast.Comment) {
+	for _, c := range comments {
+		sb.WriteString(c.Text)
+		sb.WriteString("\n")
+	}
+}
+
+// formatTrailingComments writes trailing comments (on same line) to the builder.
+func formatTrailingComments(sb *strings.Builder, comments []*ast.Comment) {
+	for _, c := range comments {
+		sb.WriteString(" ")
+		sb.WriteString(c.Text)
+	}
+}
+
 // Statement formats a single statement.
 func Statement(sb *strings.Builder, stmt ast.Statement) {
 	if stmt == nil {
 		return
 	}
 
+	// Handle statement with comments wrapper
+	if swc, ok := stmt.(*ast.StatementWithComments); ok {
+		formatComments(sb, swc.LeadingComments)
+		Statement(sb, swc.Statement)
+		formatTrailingComments(sb, swc.TrailingComments)
+		return
+	}
+
 	switch s := stmt.(type) {
 	case *ast.SelectWithUnionQuery:
 		formatSelectWithUnionQuery(sb, s)

parser/parser.go

@@ -15,10 +15,11 @@ import (
 
 // Parser parses ClickHouse SQL statements.
 type Parser struct {
-	lexer   *lexer.Lexer
-	current lexer.Item
-	peek    lexer.Item
-	errors  []error
+	lexer           *lexer.Lexer
+	current         lexer.Item
+	peek            lexer.Item
+	errors          []error
+	pendingComments []*ast.Comment // comments collected but not yet assigned to a statement
 }
 
 // New creates a new Parser from an io.Reader.
@@ -36,13 +37,28 @@ func (p *Parser) nextToken() {
 	p.current = p.peek
 	for {
 		p.peek = p.lexer.NextToken()
-		// Skip comments and whitespace
-		if p.peek.Token != token.COMMENT && p.peek.Token != token.WHITESPACE {
-			break
+		// Skip whitespace but collect comments
+		if p.peek.Token == token.WHITESPACE {
+			continue
+		}
+		if p.peek.Token == token.COMMENT {
+			p.pendingComments = append(p.pendingComments, &ast.Comment{
+				Position: p.peek.Pos,
+				Text:     p.peek.Value,
+			})
+			continue
 		}
+		break
 	}
 }
 
+// consumePendingComments returns all pending comments and clears the list.
+func (p *Parser) consumePendingComments() []*ast.Comment {
+	comments := p.pendingComments
+	p.pendingComments = nil
+	return comments
+}
+
 func (p *Parser) currentIs(t token.Token) bool {
 	return p.current.Token == t
 }
@@ -88,8 +104,22 @@ func (p *Parser) ParseStatements(ctx context.Context) ([]ast.Statement, error) {
 		default:
 		}
 
+		// Collect leading comments before the statement
+		leadingComments := p.consumePendingComments()
+
 		stmt := p.parseStatement()
 		if stmt != nil {
+			// Collect trailing comments after the statement (before semicolon or next statement)
+			trailingComments := p.consumePendingComments()
+
+			// Wrap the statement with its comments if there are any
+			if len(leadingComments) > 0 || len(trailingComments) > 0 {
+				stmt = &ast.StatementWithComments{
+					Statement:        stmt,
+					LeadingComments:  leadingComments,
+					TrailingComments: trailingComments,
+				}
+			}
 			statements = append(statements, stmt)
 		}
 

parser/parser_test.go

@@ -6,13 +6,24 @@ import (
 	"flag"
 	"os"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/sqlc-dev/doubleclick/parser"
 )
 
+// whitespaceRegex matches sequences of whitespace characters
+var whitespaceRegex = regexp.MustCompile(`\s+`)
+
+// normalizeWhitespace collapses all whitespace sequences to a single space
+// and trims leading/trailing whitespace. This allows comparing SQL statements
+// while ignoring formatting differences.
+func normalizeWhitespace(s string) string {
+	return strings.TrimSpace(whitespaceRegex.ReplaceAllString(s, " "))
+}
+
 // checkSkipped runs skipped todo tests to see which ones now pass.
 // Use with: go test ./parser -check-skipped -v
 var checkSkipped = flag.Bool("check-skipped", false, "Run skipped todo tests to see which ones now pass")
@@ -180,7 +191,10 @@ func TestParser(t *testing.T) {
 			if !metadata.TodoFormat || *checkFormat {
 				formatted := parser.Format(stmts)
 				expected := strings.TrimSpace(query)
-				if formatted != expected {
+				// Compare with whitespace normalization to ignore formatting differences
+				formattedNorm := normalizeWhitespace(formatted)
+				expectedNorm := normalizeWhitespace(expected)
+				if formattedNorm != expectedNorm {
 					if metadata.TodoFormat {
 						if *checkFormat {
 							t.Logf("FORMAT STILL FAILING:\nExpected:\n%s\n\nGot:\n%s", expected, formatted)

parser/testdata/00006_extremes_and_subquery_from/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00017_in_subquery_with_empty_result/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00024_random_counters/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00039_primary_key/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00059_shard_global_in/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00060_move_to_prewhere_and_sets/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00068_subquery_in_prewhere/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00069_date_arithmetic/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00069_duplicate_aggregation_keys/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00073_uniq_array/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00083_array_filter/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00087_where_0/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00095_hyperscan_profiler/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00097_constexpr_in_index/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00104_totals_having_mode/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00108_shard_totals_after_having/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00112_shard_totals_after_having/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00123_shard_unmerged_result_when_max_distributed_connections_is_one/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00137_in_constants/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00140_rename/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00145_aggregate_functions_statistics/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00145_empty_likes/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00151_order_by_read_in_order/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00156_max_execution_speed_sample_merge/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00171_shard_array_of_tuple_remote/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00173_compare_date_time_with_constant_string/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00176_if_string_arrays/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00183_skip_unavailable_shards/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00187_like_regexp_prefix/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00195_shard_union_all_and_global_in/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00206_empty_array_to_single/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00209_insert_select_extremes/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00213_multiple_global_in/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00224_shard_distributed_aggregation_memory_efficient_and_overflows/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00256_reverse/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00266_shard_global_subquery_and_aliases/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00270_views_query_processing_stage/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00274_shard_group_array/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00293_shard_max_subquery_depth/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00330_view_subqueries/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00337_shard_any_heavy/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00344_row_number_in_all_blocks/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00346_if_tuple/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00359_convert_or_zero_functions/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00460_vertical_and_totals_extremes/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00472_compare_uuid_with_constant_string/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00486_if_fixed_string/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00524_time_intervals_months_underflow/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00538_datediff/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00546_shard_tuple_element_formatting/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00581_limit_on_result_and_subquery_and_insert/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00612_shard_count/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00613_shard_distributed_max_execution_time/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00662_array_has_nullable/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00663_tiny_log_empty_insert/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00664_cast_from_string_to_nullable/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00676_group_by_in/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00677_shard_any_heavy_merge/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}

parser/testdata/00691_array_distinct/metadata.json

@@ -1 +1 @@
-{"todo_format":true}
+{}