From 594645b5a46c973fb4b4a6de97e20a456a563f17 Mon Sep 17 00:00:00 2001 From: Joel Marcotte Date: Wed, 30 Sep 2026 12:42:02 -0400 Subject: [PATCH 1/2] fix(normalizer): preserve CTE AS while removing aliases --- normalizer.go | 110 +++++++++++++++++++++++++++++++++++++++++++-- normalizer_test.go | 25 +++++++++++ 2 files changed, 131 insertions(+), 4 deletions(-) diff --git a/normalizer.go b/normalizer.go index a893c43..801e5b9 100644 --- a/normalizer.go +++ b/normalizer.go @@ -142,6 +142,107 @@ type headState struct { parenthesesDepth int } +type ctePhase uint8 + +const ( + cteName ctePhase = iota + cteAs + cteColumns + cteBodyStart + cteBody + cteAfterBody +) + +type cteFrame struct { + depth int + phase ctePhase +} + +type cteAliasContext struct { + depth int + frames []cteFrame + previousAS bool +} + +// isCTEAS keeps AS in a CTE header without retaining projection or table aliases. +func (c *cteAliasContext) isCTEAS(token *Token) bool { + previousAS := c.previousAS + if !isValueToken(token) { + return previousAS + } + c.previousAS = false + if token.Type == PUNCTUATION && token.Value == ";" { + c.frames = c.frames[:0] + c.depth = 0 + return false + } + depth := c.depth + if token.Type == CTE_INDICATOR { + c.frames = append(c.frames, cteFrame{depth: depth, phase: cteName}) + } else { + for len(c.frames) > 0 { + frame := &c.frames[len(c.frames)-1] + handled := false + switch frame.phase { + case cteName: + if strings.EqualFold(token.Value, "RECURSIVE") { + handled = true + } + if !handled && depth == frame.depth && (token.Type == IDENT || token.Type == QUOTED_IDENT) { + frame.phase = cteAs + handled = true + } + case cteAs: + if depth == frame.depth && token.Type == ALIAS_INDICATOR { + frame.phase = cteBodyStart + c.previousAS = true + return true + } + if depth == frame.depth && token.Type == PUNCTUATION && token.Value == "(" { + frame.phase = cteColumns + handled = true + } + case cteColumns: + if token.Type == PUNCTUATION && token.Value == ")" && depth == frame.depth+1 { + frame.phase = cteAs + } + handled = true + case cteBodyStart: + if token.Type == PUNCTUATION && token.Value == "(" && depth == frame.depth { + frame.phase = cteBody + handled = true + } + if strings.EqualFold(token.Value, "NOT") || strings.EqualFold(token.Value, "MATERIALIZED") { + handled = true + } + case cteBody: + if token.Type == PUNCTUATION && token.Value == ")" && depth == frame.depth+1 { + frame.phase = cteAfterBody + } + handled = true + case cteAfterBody: + if token.Type == PUNCTUATION && token.Value == "," && depth == frame.depth { + frame.phase = cteName + handled = true + } + if depth != frame.depth { + handled = true + } + } + if handled { + break + } + c.frames = c.frames[:len(c.frames)-1] + } + } + if token.Type == PUNCTUATION && token.Value == "(" { + c.depth++ + } else if token.Type == PUNCTUATION && token.Value == ")" && c.depth > 0 { + c.depth-- + } + return previousAS +} + type Normalizer struct { config *normalizerConfig } @@ -169,6 +270,7 @@ func (n *Normalizer) normalizeToken(lexer *Lexer, normalizedSQLBuilder *strings. var groupablePlaceholder groupablePlaceholder var headState headState var colonCtx colonContext + var cteContext cteAliasContext var ctes map[string]bool // Lazily initialized when first CTE is encountered var inTableList bool @@ -183,7 +285,7 @@ func (n *Normalizer) normalizeToken(lexer *Lexer, normalizedSQLBuilder *strings. if n.shouldCollectMetadata() { n.collectMetadata(token, lastValueToken, meta, statementMetadata, &ctes, &inTableList) } - n.normalizeSQL(token, lastValueToken, normalizedSQLBuilder, &groupablePlaceholder, &headState, &colonCtx, lexerOpts...) + n.normalizeSQL(token, lastValueToken, normalizedSQLBuilder, &groupablePlaceholder, &headState, &colonCtx, cteContext.isCTEAS(token), lexerOpts...) if token.Type == EOF { break } @@ -283,7 +385,7 @@ func (n *Normalizer) collectMetadata(token *Token, lastValueToken *LastValueToke } } -func (n *Normalizer) normalizeSQL(token *Token, lastValueToken *LastValueToken, normalizedSQLBuilder *strings.Builder, groupablePlaceholder *groupablePlaceholder, headState *headState, colonCtx *colonContext, lexerOpts ...lexerOption) { +func (n *Normalizer) normalizeSQL(token *Token, lastValueToken *LastValueToken, normalizedSQLBuilder *strings.Builder, groupablePlaceholder *groupablePlaceholder, headState *headState, colonCtx *colonContext, cteAS bool, lexerOpts ...lexerOption) { if token.Type != SPACE && token.Type != COMMENT && token.Type != MULTILINE_COMMENT { if token.Type == QUOTED_IDENT && !n.config.KeepIdentifierQuotation { if n.shouldStripIdentifierQuotes(token, lastValueToken) { @@ -338,11 +440,11 @@ func (n *Normalizer) normalizeSQL(token *Token, lastValueToken *LastValueToken, if !n.config.KeepSQLAlias { // discard SQL alias - if token.Type == ALIAS_INDICATOR { + if token.Type == ALIAS_INDICATOR && !cteAS { return } - if lastValueToken != nil && lastValueToken.Type == ALIAS_INDICATOR { + if lastValueToken != nil && lastValueToken.Type == ALIAS_INDICATOR && !cteAS { if token.Type == IDENT || token.Type == QUOTED_IDENT { return } else { diff --git a/normalizer_test.go b/normalizer_test.go index 955ae8a..3642f91 100644 --- a/normalizer_test.go +++ b/normalizer_test.go @@ -1514,6 +1514,31 @@ func TestNormalizerCTEWithoutCollectTables(t *testing.T) { } } +func TestNormalizerCTEASWithAliasRemoval(t *testing.T) { + tests := []struct { + input string + want string + }{ + {"WITH c AS MATERIALIZED (SELECT 1) SELECT * FROM c", "WITH c AS MATERIALIZED ( SELECT 1 ) SELECT * FROM c"}, + {"WITH c AS NOT MATERIALIZED (SELECT 1) SELECT * FROM c", "WITH c AS NOT MATERIALIZED ( SELECT 1 ) SELECT * FROM c"}, + {"WITH c AS (SELECT 1) SELECT * FROM c", "WITH c AS ( SELECT 1 ) SELECT * FROM c"}, + {"WITH c AS MATERIALIZED (SELECT 1 AS value) SELECT * FROM c AS alias", "WITH c AS MATERIALIZED ( SELECT 1 ) SELECT * FROM c"}, + {"WITH c AS (SELECT 1), d AS MATERIALIZED (SELECT * FROM c) SELECT * FROM d", "WITH c AS ( SELECT 1 ), d AS MATERIALIZED ( SELECT * FROM c ) SELECT * FROM d"}, + {"WITH c(v) AS (SELECT 1) SELECT * FROM c", "WITH c ( v ) AS ( SELECT 1 ) SELECT * FROM c"}, + {"WITH RECURSIVE c AS (SELECT 1) SELECT * FROM c", "WITH RECURSIVE c AS ( SELECT 1 ) SELECT * FROM c"}, + {"WITH c AS (WITH d AS (SELECT 1) SELECT * FROM d) SELECT * FROM c", "WITH c AS ( WITH d AS ( SELECT 1 ) SELECT * FROM d ) SELECT * FROM c"}, + {"SELECT 1 AS materialized FROM t AS alias", "SELECT 1 FROM t"}, + } + normalizer := NewNormalizer(WithKeepSQLAlias(false)) + for _, test := range tests { + t.Run(test.input, func(t *testing.T) { + got, _, err := normalizer.Normalize(test.input) + assert.NoError(t, err) + assert.Equal(t, test.want, got) + }) + } +} + func TestNormalizerOracleBindVariables(t *testing.T) { // This test is to ensure that Oracle bind variables with colons are normalized correctly // without adding unwanted spaces between the colon and the parameter name. From e2f9e81fb2e88fca57ef9e3496a2028d48e2dd23 Mon Sep 17 00:00:00 2001 From: Joel Marcotte Date: Wed, 30 Sep 2026 13:34:29 -0400 Subject: [PATCH 2/2] perf(normalizer): avoid CTE frame allocations --- normalizer.go | 54 ++++++++++++++++++++++++++++++------ normalizer_cte_bench_test.go | 36 ++++++++++++++++++++++++ normalizer_test.go | 5 ++++ 3 files changed, 86 insertions(+), 9 deletions(-) create mode 100644 normalizer_cte_bench_test.go diff --git a/normalizer.go b/normalizer.go index 801e5b9..f0b6361 100644 --- a/normalizer.go +++ b/normalizer.go @@ -159,9 +159,34 @@ type cteFrame struct { } type cteAliasContext struct { - depth int - frames []cteFrame - previousAS bool + depth int + frames [4]cteFrame + frameCount int + overflowFrames []cteFrame + previousAS bool +} + +func (c *cteAliasContext) pushFrame(frame cteFrame) { + if c.frameCount < len(c.frames) { + c.frames[c.frameCount] = frame + } else { + c.overflowFrames = append(c.overflowFrames, frame) + } + c.frameCount++ +} + +func (c *cteAliasContext) lastFrame() *cteFrame { + if c.frameCount <= len(c.frames) { + return &c.frames[c.frameCount-1] + } + return &c.overflowFrames[len(c.overflowFrames)-1] +} + +func (c *cteAliasContext) popFrame() { + if c.frameCount > len(c.frames) { + c.overflowFrames = c.overflowFrames[:len(c.overflowFrames)-1] + } + c.frameCount-- } // isCTEAS keeps AS in a CTE header without retaining projection or table aliases. @@ -172,16 +197,22 @@ func (c *cteAliasContext) isCTEAS(token *Token) bool { } c.previousAS = false if token.Type == PUNCTUATION && token.Value == ";" { - c.frames = c.frames[:0] + c.frameCount = 0 + c.overflowFrames = c.overflowFrames[:0] c.depth = 0 return false } depth := c.depth if token.Type == CTE_INDICATOR { - c.frames = append(c.frames, cteFrame{depth: depth, phase: cteName}) + if c.frameCount == 0 { + // Frame depths are relative to this WITH. Earlier parentheses do not matter. + c.depth = 0 + depth = 0 + } + c.pushFrame(cteFrame{depth: depth, phase: cteName}) } else { - for len(c.frames) > 0 { - frame := &c.frames[len(c.frames)-1] + for c.frameCount > 0 { + frame := c.lastFrame() handled := false switch frame.phase { case cteName: @@ -232,7 +263,7 @@ func (c *cteAliasContext) isCTEAS(token *Token) bool { if handled { break } - c.frames = c.frames[:len(c.frames)-1] + c.popFrame() } } if token.Type == PUNCTUATION && token.Value == "(" { @@ -275,6 +306,7 @@ func (n *Normalizer) normalizeToken(lexer *Lexer, normalizedSQLBuilder *strings. var inTableList bool var lastValueToken *LastValueToken + stripAliases := !n.config.KeepSQLAlias for { token := lexer.Scan() @@ -285,7 +317,11 @@ func (n *Normalizer) normalizeToken(lexer *Lexer, normalizedSQLBuilder *strings. if n.shouldCollectMetadata() { n.collectMetadata(token, lastValueToken, meta, statementMetadata, &ctes, &inTableList) } - n.normalizeSQL(token, lastValueToken, normalizedSQLBuilder, &groupablePlaceholder, &headState, &colonCtx, cteContext.isCTEAS(token), lexerOpts...) + cteAS := false + if stripAliases && (cteContext.frameCount > 0 || token.Type == CTE_INDICATOR) { + cteAS = cteContext.isCTEAS(token) + } + n.normalizeSQL(token, lastValueToken, normalizedSQLBuilder, &groupablePlaceholder, &headState, &colonCtx, cteAS, lexerOpts...) if token.Type == EOF { break } diff --git a/normalizer_cte_bench_test.go b/normalizer_cte_bench_test.go new file mode 100644 index 0000000..039d21b --- /dev/null +++ b/normalizer_cte_bench_test.go @@ -0,0 +1,36 @@ +package sqllexer + +import ( + "strings" + "testing" +) + +var normalizedCTEBenchmarkResult string + +func BenchmarkCTENormalization(b *testing.B) { + queries := []struct { + name string + query string + }{ + {"NoCTE", "SELECT id AS identifier FROM users AS u WHERE id = 42"}, + {"Plain", "WITH c AS (SELECT 1) SELECT * FROM c"}, + {"Materialized", "WITH c AS MATERIALIZED (SELECT 1 AS value) SELECT * FROM c AS alias"}, + {"NotMaterialized", "WITH c AS NOT MATERIALIZED (SELECT 1) SELECT * FROM c"}, + {"Nested", "WITH c AS (WITH d AS (SELECT 1) SELECT * FROM d) SELECT * FROM c"}, + {"DeepNested", strings.Repeat("WITH c AS (", 5) + "SELECT 1" + strings.Repeat(") SELECT * FROM c", 5)}, + } + normalizer := NewNormalizer(WithKeepSQLAlias(false)) + for _, benchmark := range queries { + b.Run(benchmark.name, func(b *testing.B) { + b.ReportAllocs() + b.SetBytes(int64(len(benchmark.query))) + for i := 0; i < b.N; i++ { + output, _, err := normalizer.Normalize(benchmark.query) + if err != nil { + b.Fatal(err) + } + normalizedCTEBenchmarkResult = output + } + }) + } +} diff --git a/normalizer_test.go b/normalizer_test.go index 3642f91..9109ffb 100644 --- a/normalizer_test.go +++ b/normalizer_test.go @@ -2,6 +2,7 @@ package sqllexer import ( "fmt" + "strings" "testing" "github.com/stretchr/testify/assert" @@ -1527,6 +1528,10 @@ func TestNormalizerCTEASWithAliasRemoval(t *testing.T) { {"WITH c(v) AS (SELECT 1) SELECT * FROM c", "WITH c ( v ) AS ( SELECT 1 ) SELECT * FROM c"}, {"WITH RECURSIVE c AS (SELECT 1) SELECT * FROM c", "WITH RECURSIVE c AS ( SELECT 1 ) SELECT * FROM c"}, {"WITH c AS (WITH d AS (SELECT 1) SELECT * FROM d) SELECT * FROM c", "WITH c AS ( WITH d AS ( SELECT 1 ) SELECT * FROM d ) SELECT * FROM c"}, + {strings.Repeat("WITH c AS (", 5) + "SELECT 1" + strings.Repeat(") SELECT * FROM c", 5), strings.Repeat("WITH c AS ( ", 5) + "SELECT 1" + strings.Repeat(" ) SELECT * FROM c", 5)}, + {"SELECT * FROM (WITH c AS (SELECT 1) SELECT * FROM c) AS t", "SELECT * FROM ( WITH c AS ( SELECT 1 ) SELECT * FROM c )"}, + {"SELECT * FROM (WITH c AS (SELECT 1) SELECT * FROM c) AS t UNION ALL SELECT * FROM (WITH d AS (SELECT 2) SELECT * FROM d) AS u", "SELECT * FROM ( WITH c AS ( SELECT 1 ) SELECT * FROM c ) UNION ALL SELECT * FROM ( WITH d AS ( SELECT 2 ) SELECT * FROM d )"}, + {"wItH c aS (SELECT 1) SELECT * FROM c", "wItH c aS ( SELECT 1 ) SELECT * FROM c"}, {"SELECT 1 AS materialized FROM t AS alias", "SELECT 1 FROM t"}, } normalizer := NewNormalizer(WithKeepSQLAlias(false))