From fb16511a7d5830feb90226262e879e05cdbae2ea Mon Sep 17 00:00:00 2001 From: Joel Marcotte Date: Wed, 30 Sep 2026 14:14:27 -0400 Subject: [PATCH] Preserve AS in CAST expressions during normalization --- normalizer.go | 110 +++++++++++++++++++++++++++++++--- normalizer_cast_bench_test.go | 32 ++++++++++ normalizer_cast_test.go | 48 +++++++++++++++ 3 files changed, 183 insertions(+), 7 deletions(-) create mode 100644 normalizer_cast_bench_test.go create mode 100644 normalizer_cast_test.go diff --git a/normalizer.go b/normalizer.go index f0b6361..388d891 100644 --- a/normalizer.go +++ b/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 { diff --git a/normalizer_cast_bench_test.go b/normalizer_cast_bench_test.go new file mode 100644 index 0000000..b53f1d1 --- /dev/null +++ b/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 + } + }) + } +} diff --git a/normalizer_cast_test.go b/normalizer_cast_test.go new file mode 100644 index 0000000..6e16ab8 --- /dev/null +++ b/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) + } +}