From eea9029022afecba09736be74bb6e1ab4916aab6 Mon Sep 17 00:00:00 2001 From: Joel Marcotte Date: Tue, 5 May 2026 10:56:15 -0400 Subject: [PATCH] [SDBM-2597] Obfuscate EXTRACT field keywords The pg_stat_statements form `EXTRACT($1 FROM x)` and the pg_stat_activity form `EXTRACT(epoch FROM x)` previously produced different obfuscated signatures: the positional parameter was replaced with `?` while the unquoted field keyword was preserved as an identifier. This split one logical query into two DBM signatures. Track EXTRACT( context in the obfuscator and replace the field argument with `?` when it matches a known PostgreSQL field name (epoch, year, month, dow, etc.). Bare `epoch` outside an EXTRACT call and unrecognized field names are left untouched. Co-Authored-By: Claude Opus 4.7 (1M context) --- obfuscate_and_normalize.go | 3 + obfuscator.go | 4 + obfuscator_test.go | 28 +++++++ sqllexer_utils.go | 78 +++++++++++++++++++ .../select/date-part-string-literal.json | 8 ++ .../select/date-trunc-string-literal.json | 8 ++ .../select/extract-epoch-field.json | 26 +++++++ .../select/extract-field-name-as-column.json | 8 ++ .../select/extract-positional-parameter.json | 8 ++ .../select/extract-quoted-string-field.json | 8 ++ .../select/extract-various-fields.json | 8 ++ 11 files changed, 187 insertions(+) create mode 100644 testdata/postgresql/select/date-part-string-literal.json create mode 100644 testdata/postgresql/select/date-trunc-string-literal.json create mode 100644 testdata/postgresql/select/extract-epoch-field.json create mode 100644 testdata/postgresql/select/extract-field-name-as-column.json create mode 100644 testdata/postgresql/select/extract-positional-parameter.json create mode 100644 testdata/postgresql/select/extract-quoted-string-field.json create mode 100644 testdata/postgresql/select/extract-various-fields.json diff --git a/obfuscate_and_normalize.go b/obfuscate_and_normalize.go index 1e513cd..81857dd 100644 --- a/obfuscate_and_normalize.go +++ b/obfuscate_and_normalize.go @@ -3,8 +3,11 @@ 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 ec extractContext obfuscate := func(token *Token, lastValueToken *LastValueToken) { obfuscator.ObfuscateTokenValue(token, lastValueToken, lexerOpts...) + ec.maybeReplaceExtractField(token) + ec.update(token) } return normalizer.normalize(input, obfuscate, lexerOpts...) } diff --git a/obfuscator.go b/obfuscator.go index 1bb91c8..4c00947 100644 --- a/obfuscator.go +++ b/obfuscator.go @@ -91,6 +91,7 @@ func (o *Obfuscator) Obfuscate(input string, lexerOpts ...lexerOption) string { ) var lastValueToken *LastValueToken + var ec extractContext for { token := lexer.Scan() @@ -98,10 +99,13 @@ func (o *Obfuscator) Obfuscate(input string, lexerOpts ...lexerOption) string { break } o.ObfuscateTokenValue(token, lastValueToken, lexerOpts...) + ec.maybeReplaceExtractField(token) + obfuscatedSQL.WriteString(token.Value) if isValueToken(token) { lastValueToken = token.getLastValueToken() } + ec.update(token) } return strings.Clone(strings.TrimSpace(obfuscatedSQL.String())) diff --git a/obfuscator_test.go b/obfuscator_test.go index f384a8b..df3dbf0 100644 --- a/obfuscator_test.go +++ b/obfuscator_test.go @@ -551,6 +551,34 @@ func TestObfuscator(t *testing.T) { expected: `SELECT * FROM users where id = ?`, replaceBindParameter: true, }, + { + // pg_stat_activity captures the raw SQL with `epoch` as an unquoted + // identifier. Obfuscate it so the signature converges with the + // pg_stat_statements form `EXTRACT($1 FROM ...)`. + input: `SELECT EXTRACT(epoch FROM created_at) FROM events`, + expected: `SELECT EXTRACT(? FROM created_at) FROM events`, + }, + { + input: `SELECT EXTRACT(YEAR FROM a), EXTRACT(MoNtH FROM b) FROM t`, + expected: `SELECT EXTRACT(? FROM a), EXTRACT(? FROM b) FROM t`, + }, + { + input: `SELECT EXTRACT($1 FROM created_at) FROM events`, + expected: `SELECT EXTRACT(? FROM created_at) FROM events`, + replacePositionalParameter: true, + }, + { + // `epoch` outside an EXTRACT(...) call is just an identifier and + // must not be replaced. + input: `SELECT epoch FROM events WHERE epoch > 0`, + expected: `SELECT epoch FROM events WHERE epoch > ?`, + }, + { + // Unknown EXTRACT field (not in the recognized set) is left alone + // so we don't accidentally rewrite user identifiers. + input: `SELECT EXTRACT(custom_field FROM created_at) FROM events`, + expected: `SELECT EXTRACT(custom_field FROM created_at) FROM events`, + }, } for _, tt := range tests { diff --git a/sqllexer_utils.go b/sqllexer_utils.go index ab717de..b36f2f5 100644 --- a/sqllexer_utils.go +++ b/sqllexer_utils.go @@ -179,8 +179,86 @@ var ( alias = []string{ "AS", } + + // extractFieldKeywords is the set of PostgreSQL EXTRACT field names + // (https://www.postgresql.org/docs/current/functions-datetime.html#FUNCTIONS-DATETIME-EXTRACT). + // They appear as unquoted identifiers in `EXTRACT(field FROM source)` and must be + // obfuscated to a placeholder so signatures match the pg_stat_statements form + // `EXTRACT($1 FROM source)`. + extractFieldKeywords = map[string]struct{}{ + "CENTURY": {}, + "DAY": {}, + "DECADE": {}, + "DOW": {}, + "DOY": {}, + "EPOCH": {}, + "HOUR": {}, + "ISODOW": {}, + "ISOYEAR": {}, + "JULIAN": {}, + "MICROSECONDS": {}, + "MILLENNIUM": {}, + "MILLISECONDS": {}, + "MINUTE": {}, + "MONTH": {}, + "QUARTER": {}, + "SECOND": {}, + "TIMEZONE": {}, + "TIMEZONE_HOUR": {}, + "TIMEZONE_MINUTE": {}, + "WEEK": {}, + "YEAR": {}, + } ) +// isExtractFieldKeyword reports whether value is a known PostgreSQL EXTRACT field +// keyword (case-insensitive). +func isExtractFieldKeyword(value string) bool { + _, ok := extractFieldKeywords[strings.ToUpper(value)] + return ok +} + +// extractContext tracks the two most recent non-whitespace, non-comment tokens +// so callers can detect when the current token is the field argument of an +// `EXTRACT(field FROM ...)` call. The zero value is ready to use. +type extractContext struct { + prevType TokenType + prevValue string + prev2Type TokenType + prev2Value string +} + +// maybeReplaceExtractField rewrites token.Value to StringPlaceholder when the +// token is the field argument of an EXTRACT(...) call (i.e. preceded by +// FUNCTION "EXTRACT" then PUNCTUATION "(") and matches a known field name. +// This converges the obfuscated form of `EXTRACT(epoch FROM x)` with the +// pg_stat_statements form `EXTRACT($1 FROM x)`. +func (c *extractContext) maybeReplaceExtractField(token *Token) { + if token.Type != IDENT && token.Type != KEYWORD { + return + } + if c.prevType != PUNCTUATION || c.prevValue != "(" { + return + } + if c.prev2Type != FUNCTION || !strings.EqualFold(c.prev2Value, "EXTRACT") { + return + } + if !isExtractFieldKeyword(token.Value) { + return + } + token.Value = StringPlaceholder +} + +// update advances the rolling window of recent tokens, ignoring whitespace and +// comments so they don't break the EXTRACT( ... ) adjacency check. +func (c *extractContext) update(token *Token) { + if token.Type == SPACE || token.Type == COMMENT || token.Type == MULTILINE_COMMENT { + return + } + c.prev2Type, c.prev2Value = c.prevType, c.prevValue + c.prevType, c.prevValue = token.Type, token.Value +} + // trieNode represents a node in the keyword trie. type trieNode struct { children [27]*trieNode // 0-25 for A-Z, 26 for underscore diff --git a/testdata/postgresql/select/date-part-string-literal.json b/testdata/postgresql/select/date-part-string-literal.json new file mode 100644 index 0000000..73c9e87 --- /dev/null +++ b/testdata/postgresql/select/date-part-string-literal.json @@ -0,0 +1,8 @@ +{ + "input": "SELECT DATE_PART('hour', created_at) FROM events;", + "outputs": [ + { + "expected": "SELECT DATE_PART ( ?, created_at ) FROM events" + } + ] +} diff --git a/testdata/postgresql/select/date-trunc-string-literal.json b/testdata/postgresql/select/date-trunc-string-literal.json new file mode 100644 index 0000000..9756346 --- /dev/null +++ b/testdata/postgresql/select/date-trunc-string-literal.json @@ -0,0 +1,8 @@ +{ + "input": "SELECT DATE_TRUNC('month', created_at) FROM events;", + "outputs": [ + { + "expected": "SELECT DATE_TRUNC ( ?, created_at ) FROM events" + } + ] +} diff --git a/testdata/postgresql/select/extract-epoch-field.json b/testdata/postgresql/select/extract-epoch-field.json new file mode 100644 index 0000000..c0beff6 --- /dev/null +++ b/testdata/postgresql/select/extract-epoch-field.json @@ -0,0 +1,26 @@ +{ + "input": "SELECT EXTRACT(epoch FROM created_at) FROM events;", + "outputs": [ + { + "expected": "SELECT EXTRACT ( ? FROM created_at ) FROM events", + "statement_metadata": { + "size": 22, + "tables": [ + "created_at", + "events" + ], + "commands": [ + "SELECT" + ], + "comments": [], + "procedures": [] + } + }, + { + "expected": "SELECT EXTRACT(? FROM created_at) FROM events", + "normalizer_config": { + "remove_space_between_parentheses": true + } + } + ] +} diff --git a/testdata/postgresql/select/extract-field-name-as-column.json b/testdata/postgresql/select/extract-field-name-as-column.json new file mode 100644 index 0000000..06e6e0f --- /dev/null +++ b/testdata/postgresql/select/extract-field-name-as-column.json @@ -0,0 +1,8 @@ +{ + "input": "SELECT epoch, year FROM events WHERE epoch > 100;", + "outputs": [ + { + "expected": "SELECT epoch, year FROM events WHERE epoch > ?" + } + ] +} diff --git a/testdata/postgresql/select/extract-positional-parameter.json b/testdata/postgresql/select/extract-positional-parameter.json new file mode 100644 index 0000000..8097cc2 --- /dev/null +++ b/testdata/postgresql/select/extract-positional-parameter.json @@ -0,0 +1,8 @@ +{ + "input": "SELECT EXTRACT($1 FROM created_at) FROM events;", + "outputs": [ + { + "expected": "SELECT EXTRACT ( ? FROM created_at ) FROM events" + } + ] +} diff --git a/testdata/postgresql/select/extract-quoted-string-field.json b/testdata/postgresql/select/extract-quoted-string-field.json new file mode 100644 index 0000000..0da4463 --- /dev/null +++ b/testdata/postgresql/select/extract-quoted-string-field.json @@ -0,0 +1,8 @@ +{ + "input": "SELECT EXTRACT('epoch' FROM created_at) FROM events;", + "outputs": [ + { + "expected": "SELECT EXTRACT ( ? FROM created_at ) FROM events" + } + ] +} diff --git a/testdata/postgresql/select/extract-various-fields.json b/testdata/postgresql/select/extract-various-fields.json new file mode 100644 index 0000000..c19562b --- /dev/null +++ b/testdata/postgresql/select/extract-various-fields.json @@ -0,0 +1,8 @@ +{ + "input": "SELECT EXTRACT(year FROM a), EXTRACT(MONTH FROM b), EXTRACT(dow FROM c), EXTRACT(microseconds FROM d), EXTRACT(timezone_hour FROM e) FROM t;", + "outputs": [ + { + "expected": "SELECT EXTRACT ( ? FROM a ), EXTRACT ( ? FROM b ), EXTRACT ( ? FROM c ), EXTRACT ( ? FROM d ), EXTRACT ( ? FROM e ) FROM t" + } + ] +}