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" + } + ] +}