Fix CTE AS normalization with alias removal

normalizer.go

@@ -142,6 +142,138 @@ 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         [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.
+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.frameCount = 0
+		c.overflowFrames = c.overflowFrames[:0]
+		c.depth = 0
+		return false
+	}
+	depth := c.depth
+	if token.Type == CTE_INDICATOR {
+		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 c.frameCount > 0 {
+			frame := c.lastFrame()
+			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.popFrame()
+		}
+	}
+	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,10 +301,12 @@ 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
 
 	var lastValueToken *LastValueToken
+	stripAliases := !n.config.KeepSQLAlias
 
 	for {
 		token := lexer.Scan()
@@ -183,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, 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
 		}
@@ -283,7 +421,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 +476,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 {

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
+			}
+		})
+	}
+}

normalizer_test.go

@@ -2,6 +2,7 @@ package sqllexer
 
 import (
 	"fmt"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -1514,6 +1515,35 @@ 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"},
+		{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))
+	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.