Preserve AS in CAST expressions during normalization

normalizer.go

@@ -274,6 +274,93 @@ func (c *cteAliasContext) isCTEAS(token *Token) bool {
 	return previousAS
 }
 
+type castAliasContext struct {
+	depth          int
+	frames         [4]int
+	frameCount     int
+	overflowFrames []int
+	pending        bool
+	previousAS     bool
+}
+
+func (c *castAliasContext) pushFrame(depth int) {
+	if c.frameCount < len(c.frames) {
+		c.frames[c.frameCount] = depth
+	} else {
+		c.overflowFrames = append(c.overflowFrames, depth)
+	}
+	c.frameCount++
+}
+
+func (c *castAliasContext) lastFrame() int {
+	if c.frameCount <= len(c.frames) {
+		return c.frames[c.frameCount-1]
+	}
+	return c.overflowFrames[len(c.overflowFrames)-1]
+}
+
+func (c *castAliasContext) popFrame() {
+	if c.frameCount > len(c.frames) {
+		c.overflowFrames = c.overflowFrames[:len(c.overflowFrames)-1]
+	}
+	c.frameCount--
+}
+
+func isCastToken(token *Token) bool {
+	if (token.Type != FUNCTION && token.Type != IDENT) || len(token.Value) != 4 {
+		return false
+	}
+	v := token.Value
+	return v[0]|0x20 == 'c' && v[1]|0x20 == 'a' && v[2]|0x20 == 's' && v[3]|0x20 == 't'
+}
+
+// isCastAS keeps the type separator in CAST(expression AS type), including nested casts.
+func (c *castAliasContext) isCastAS(token *Token) bool {
+	previousAS := c.previousAS
+	if !isValueToken(token) {
+		return previousAS
+	}
+	c.previousAS = false
+	if token.Type == PUNCTUATION && token.Value == ";" {
+		c.depth = 0
+		c.frameCount = 0
+		c.overflowFrames = c.overflowFrames[:0]
+		c.pending = false
+		return false
+	}
+	if c.pending {
+		c.pending = false
+		if token.Type == PUNCTUATION && token.Value == "(" {
+			c.pushFrame(c.depth)
+		}
+	}
+	if c.frameCount > 0 && token.Type == ALIAS_INDICATOR && c.depth == c.lastFrame()+1 {
+		c.previousAS = true
+		return true
+	}
+	if isCastToken(token) {
+		if c.frameCount == 0 {
+			// Frame depths are relative to this CAST.
+			c.depth = 0
+		}
+		c.pending = true
+	}
+	if token.Type == PUNCTUATION {
+		switch token.Value {
+		case "(":
+			c.depth++
+		case ")":
+			if c.frameCount > 0 && c.depth == c.lastFrame()+1 {
+				c.popFrame()
+			}
+			if c.depth > 0 {
+				c.depth--
+			}
+		}
+	}
+	return previousAS
+}
+
 type Normalizer struct {
 	config *normalizerConfig
 }
@@ -302,6 +389,7 @@ func (n *Normalizer) normalizeToken(lexer *Lexer, normalizedSQLBuilder *strings.
 	var headState headState
 	var colonCtx colonContext
 	var cteContext cteAliasContext
+	var castContext castAliasContext
 	var ctes map[string]bool // Lazily initialized when first CTE is encountered
 	var inTableList bool
 
@@ -317,11 +405,14 @@ func (n *Normalizer) normalizeToken(lexer *Lexer, normalizedSQLBuilder *strings.
 		if n.shouldCollectMetadata() {
 			n.collectMetadata(token, lastValueToken, meta, statementMetadata, &ctes, &inTableList)
 		}
-		cteAS := false
+		keepAS := false
 		if stripAliases && (cteContext.frameCount > 0 || token.Type == CTE_INDICATOR) {
-			cteAS = cteContext.isCTEAS(token)
+			keepAS = cteContext.isCTEAS(token)
+		}
+		if stripAliases && (castContext.frameCount > 0 || castContext.pending || isCastToken(token)) {
+			keepAS = castContext.isCastAS(token) || keepAS
 		}
-		n.normalizeSQL(token, lastValueToken, normalizedSQLBuilder, &groupablePlaceholder, &headState, &colonCtx, cteAS, lexerOpts...)
+		n.normalizeSQL(token, lastValueToken, normalizedSQLBuilder, &groupablePlaceholder, &headState, &colonCtx, keepAS, lexerOpts...)
 		if token.Type == EOF {
 			break
 		}
@@ -342,7 +433,12 @@ func (n *Normalizer) Normalize(input string, lexerOpts ...lexerOption) (normaliz
 func (n *Normalizer) normalize(input string, preProcessToken func(*Token, *LastValueToken), lexerOpts ...lexerOption) (normalizedSQL string, statementMetadata *StatementMetadata, err error) {
 	lexer := New(input, lexerOpts...)
 	var normalizedSQLBuilder strings.Builder
-	normalizedSQLBuilder.Grow(len(input))
+	bufferSize := len(input)
+	if bufferSize <= 40 {
+		// Short statements can grow when punctuation gains surrounding spaces.
+		bufferSize += 8
+	}
+	normalizedSQLBuilder.Grow(bufferSize)
 
 	meta := &metadataSet{
 		tablesSet:     map[string]struct{}{},
@@ -421,7 +517,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, cteAS bool, lexerOpts ...lexerOption) {
+func (n *Normalizer) normalizeSQL(token *Token, lastValueToken *LastValueToken, normalizedSQLBuilder *strings.Builder, groupablePlaceholder *groupablePlaceholder, headState *headState, colonCtx *colonContext, keepAS 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) {
@@ -476,11 +572,11 @@ func (n *Normalizer) normalizeSQL(token *Token, lastValueToken *LastValueToken,
 
 		if !n.config.KeepSQLAlias {
 			// discard SQL alias
-			if token.Type == ALIAS_INDICATOR && !cteAS {
+			if token.Type == ALIAS_INDICATOR && !keepAS {
 				return
 			}
 
-			if lastValueToken != nil && lastValueToken.Type == ALIAS_INDICATOR && !cteAS {
+			if lastValueToken != nil && lastValueToken.Type == ALIAS_INDICATOR && !keepAS {
 				if token.Type == IDENT || token.Type == QUOTED_IDENT {
 					return
 				} else {

normalizer_cast_bench_test.go

@@ -0,0 +1,32 @@
+package sqllexer
+
+import "testing"
+
+var normalizedCastBenchmarkResult string
+
+func BenchmarkCastNormalization(b *testing.B) {
+	queries := []struct {
+		name  string
+		query string
+	}{
+		{"NoCast", "SELECT id AS identifier FROM users AS u WHERE id = ?"},
+		{"Simple", "SELECT CAST(? AS integer) FROM t"},
+		{"Nested", "SELECT CAST(CAST(? AS text) AS varchar(10)) FROM t"},
+		{"Expression", "SELECT CAST(COALESCE(x, ?) + ? AS numeric(10,2)) AS amount FROM t AS alias"},
+		{"CTEWithCast", "WITH c AS (SELECT CAST(? AS integer) AS value) SELECT * FROM c AS alias"},
+	}
+	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, WithDBMS(DBMSPostgres))
+				if err != nil {
+					b.Fatal(err)
+				}
+				normalizedCastBenchmarkResult = output
+			}
+		})
+	}
+}

normalizer_cast_test.go

@@ -0,0 +1,48 @@
+package sqllexer
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestNormalizerCastASAndAliases(t *testing.T) {
+	tests := []struct {
+		name  string
+		input string
+		want  string
+	}{
+		{"simple", "SELECT CAST(1 AS integer)", "SELECT CAST ( 1 AS integer )"},
+		{"spaced and mixed case", "SELECT cAsT (1 aS integer)", "SELECT cAsT ( 1 aS integer )"},
+		{"nested casts", "SELECT CAST(CAST(1 AS text) AS varchar(10))", "SELECT CAST ( CAST ( 1 AS text ) AS varchar ( 10 ) )"},
+		{"deeply nested casts", "SELECT " + strings.Repeat("CAST(", 5) + "1" + strings.Repeat(" AS integer)", 5), "SELECT " + strings.Repeat("CAST ( ", 5) + "1" + strings.Repeat(" AS integer )", 5)},
+		{"expression with nested alias", "SELECT CAST((SELECT x AS y FROM t) AS integer)", "SELECT CAST ( ( SELECT x FROM t ) AS integer )"},
+		{"projection and table aliases", "SELECT CAST(1 AS integer) AS amount FROM t AS alias", "SELECT CAST ( 1 AS integer ) FROM t"},
+		{"cast in CTE", "WITH c AS (SELECT CAST(1 AS integer) AS value) SELECT * FROM c AS alias", "WITH c AS ( SELECT CAST ( 1 AS integer ) ) SELECT * FROM c"},
+		{"ordinary aliases", "SELECT x AS amount FROM t AS alias", "SELECT x FROM t"},
+		{"cast as alias", "SELECT x AS cast FROM t", "SELECT x FROM t"},
+	}
+	normalizer := NewNormalizer(WithKeepSQLAlias(false))
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got, _, err := normalizer.Normalize(test.input, WithDBMS(DBMSPostgres))
+			if err != nil {
+				t.Fatal(err)
+			}
+			if got != test.want {
+				t.Errorf("Normalize(%q) = %q, want %q", test.input, got, test.want)
+			}
+		})
+	}
+}
+
+func TestNormalizerCastKeepAliases(t *testing.T) {
+	input := "SELECT CAST(1 AS integer) AS amount FROM t AS alias"
+	got, _, err := NewNormalizer(WithKeepSQLAlias(true)).Normalize(input, WithDBMS(DBMSPostgres))
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := "SELECT CAST ( 1 AS integer ) AS amount FROM t AS alias"
+	if got != want {
+		t.Errorf("Normalize(%q) = %q, want %q", input, got, want)
+	}
+}