From fdbe07e52a45034b4bd18921b5b641ace603c7cf Mon Sep 17 00:00:00 2001 From: Oscar Le Dauphin Date: Mon, 4 May 2026 15:57:36 +0200 Subject: [PATCH] fix(obfuscator): stop obfuscating postgres EXTRACT/date_part/date_trunc temporal field --- obfuscate_and_normalize.go | 6 +++ obfuscator.go | 65 +++++++++++++++++++++++++++++++++ obfuscator_test.go | 75 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 146 insertions(+) diff --git a/obfuscate_and_normalize.go b/obfuscate_and_normalize.go index 1e513cd..34ea124 100644 --- a/obfuscate_and_normalize.go +++ b/obfuscate_and_normalize.go @@ -3,7 +3,13 @@ package sqllexer // ObfuscateAndNormalize takes an input SQL string and returns an normalized SQL string with metadata // This function is a convenience function that combines the Obfuscator and Normalizer in one pass func ObfuscateAndNormalize(input string, obfuscator *Obfuscator, normalizer *Normalizer, lexerOpts ...lexerOption) (normalizedSQL string, statementMetadata *StatementMetadata, err error) { + var tracker temporalFuncTracker obfuscate := func(token *Token, lastValueToken *LastValueToken) { + // Preserve quoted field-name arguments (e.g. 'epoch') inside temporal + // functions so the signature matches pg_stat_statements normalization. + if tracker.advance(token) { + return + } obfuscator.ObfuscateTokenValue(token, lastValueToken, lexerOpts...) } return normalizer.normalize(input, obfuscate, lexerOpts...) diff --git a/obfuscator.go b/obfuscator.go index 1bb91c8..b8bdd43 100644 --- a/obfuscator.go +++ b/obfuscator.go @@ -79,6 +79,58 @@ const ( NumberPlaceholder = "?" ) +// temporalFunctions is the set of SQL functions whose first argument is a field +// keyword (e.g. 'epoch', 'month') rather than a user-supplied value. +var temporalFunctions = map[string]bool{ + "extract": true, "date_part": true, "date_trunc": true, +} + +func isTemporalFunction(name string) bool { + return temporalFunctions[strings.ToLower(name)] +} + +type temporalState int + +const ( + temporalStateNone temporalState = iota + temporalStateFunc // saw a temporal function token + temporalStateOpen // saw '(' after temporal function; next STRING is the field name +) + +// temporalFuncTracker detects when the current token is a quoted field-name +// argument inside a temporal function call (EXTRACT, date_part, date_trunc). +// Call advance before deciding whether to obfuscate a token; if it returns +// true the token is a field name and its quotes have already been stripped. +type temporalFuncTracker struct { + state temporalState +} + +func (t *temporalFuncTracker) advance(token *Token) (isFieldName bool) { + if token.Type == SPACE { + return false + } + switch t.state { + case temporalStateFunc: + if token.Type == PUNCTUATION && token.Value == "(" { + t.state = temporalStateOpen + } else { + t.state = temporalStateNone + } + case temporalStateOpen: + t.state = temporalStateNone + if token.Type == STRING && len(token.Value) >= 2 { + // Strip surrounding single quotes: 'epoch' → epoch + token.Value = token.Value[1 : len(token.Value)-1] + return true + } + } + // Re-evaluate as None (handles the case where state just reset above on the same token) + if t.state == temporalStateNone && token.Type == FUNCTION && isTemporalFunction(token.Value) { + t.state = temporalStateFunc + } + return false +} + // Obfuscate takes an input SQL string and returns an obfuscated SQL string. // The obfuscator replaces all literal values with a single placeholder func (o *Obfuscator) Obfuscate(input string, lexerOpts ...lexerOption) string { @@ -91,12 +143,25 @@ func (o *Obfuscator) Obfuscate(input string, lexerOpts ...lexerOption) string { ) var lastValueToken *LastValueToken + var tracker temporalFuncTracker for { token := lexer.Scan() if token.Type == EOF { break } + + // Preserve quoted field-name arguments inside temporal functions: + // EXTRACT('epoch' FROM x) and date_part('month', x) use string literals + // that are SQL keywords, not user values — they must not be replaced with ?. + if tracker.advance(token) { + obfuscatedSQL.WriteString(token.Value) + if isValueToken(token) { + lastValueToken = token.getLastValueToken() + } + continue + } + o.ObfuscateTokenValue(token, lastValueToken, lexerOpts...) obfuscatedSQL.WriteString(token.Value) if isValueToken(token) { diff --git a/obfuscator_test.go b/obfuscator_test.go index f384a8b..ad8adc2 100644 --- a/obfuscator_test.go +++ b/obfuscator_test.go @@ -579,6 +579,81 @@ func ExampleObfuscator() { // TestObfuscatorDoesNotPinLargeBackingArrays verifies that the Obfuscate function // returns a string that doesn't hold a reference to an excessively large backing array. +// TestObfuscatorTemporalFunctions verifies that quoted field-name arguments inside +// EXTRACT / date_part / date_trunc are preserved rather than replaced with ?. +// See https://github.com/DataDog/datadog-agent/issues/49420 +func TestObfuscatorTemporalFunctions(t *testing.T) { + obfuscator := NewObfuscator(WithReplacePositionalParameter(true)) + tests := []struct { + input string + expected string + }{ + // Unquoted form — baseline, must stay unchanged. + { + `SELECT EXTRACT(epoch FROM t.ts) FROM t WHERE id = 1`, + `SELECT EXTRACT(epoch FROM t.ts) FROM t WHERE id = ?`, + }, + // Quoted form (PG ≤13 / some ORMs): 'epoch' must be preserved, not replaced with ?. + { + `SELECT EXTRACT('epoch' FROM t.ts) FROM t WHERE id = 1`, + `SELECT EXTRACT(epoch FROM t.ts) FROM t WHERE id = ?`, + }, + // pg_stat_statements sends positional params; output must match the quoted-form output. + { + `SELECT EXTRACT(epoch FROM t.ts) FROM t WHERE id = $1`, + `SELECT EXTRACT(epoch FROM t.ts) FROM t WHERE id = ?`, + }, + { + `SELECT date_part('month', t.ts) FROM t WHERE id = 1`, + `SELECT date_part(month, t.ts) FROM t WHERE id = ?`, + }, + { + `SELECT date_trunc('year', t.ts) FROM t WHERE id = 1`, + `SELECT date_trunc(year, t.ts) FROM t WHERE id = ?`, + }, + // Other string args must still be obfuscated. + { + `SELECT f('secret') FROM t`, + `SELECT f(?) FROM t`, + }, + } + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + assert.Equal(t, tt.expected, obfuscator.Obfuscate(tt.input, WithDBMS(DBMSPostgres))) + }) + } +} + +func TestObfuscateAndNormalizeTemporalFunctions(t *testing.T) { + obfuscator := NewObfuscator(WithReplacePositionalParameter(true)) + normalizer := NewNormalizer() + tests := []struct { + input string + expected string + }{ + { + `SELECT EXTRACT('epoch' FROM t.ts) FROM t WHERE id = 1`, + `SELECT EXTRACT ( epoch FROM t.ts ) FROM t WHERE id = ?`, + }, + // Both forms must collapse to the same signature. + { + `SELECT EXTRACT(epoch FROM t.ts) FROM t WHERE id = $1`, + `SELECT EXTRACT ( epoch FROM t.ts ) FROM t WHERE id = ?`, + }, + { + `SELECT date_part('month', t.ts) FROM t WHERE id = 1`, + `SELECT date_part ( month, t.ts ) FROM t WHERE id = ?`, + }, + } + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + got, _, err := ObfuscateAndNormalize(tt.input, obfuscator, normalizer, WithDBMS(DBMSPostgres)) + assert.NoError(t, err) + assert.Equal(t, tt.expected, got) + }) + } +} + func TestObfuscatorDoesNotPinLargeBackingArrays(t *testing.T) { obfuscator := NewObfuscator()