Obfuscate EXTRACT field keywords

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

obfuscator.go

@@ -91,17 +91,21 @@ func (o *Obfuscator) Obfuscate(input string, lexerOpts ...lexerOption) string {
 	)
 
 	var lastValueToken *LastValueToken
+	var ec extractContext
 
 	for {
 		token := lexer.Scan()
 		if token.Type == EOF {
 			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()))

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 {

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

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

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

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

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

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

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

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