fix(obfuscator): stop obfuscating postgres EXTRACT temporal field

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...)

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) {

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()