From 2e56ef0b4ec00dcefc2c2328b7e162a8ccfcee8c Mon Sep 17 00:00:00 2001 From: Joel Marcotte Date: Wed, 30 Sep 2026 14:25:28 -0400 Subject: [PATCH 1/3] Add opt-in PostgreSQL IS predicate literal preservation --- README.md | 18 ++++ cmd/sqllexer/main.go | 5 ++ dbms_test.go | 1 + is_predicate.go | 64 +++++++++++++ is_predicate_test.go | 179 +++++++++++++++++++++++++++++++++++++ obfuscate_and_normalize.go | 13 ++- obfuscator.go | 15 +++- 7 files changed, 293 insertions(+), 2 deletions(-) create mode 100644 is_predicate.go create mode 100644 is_predicate_test.go diff --git a/README.md b/README.md index a9e0c24..a8f5f29 100644 --- a/README.md +++ b/README.md @@ -114,8 +114,26 @@ sqllexer -mode tokenize -input query.sql # Obfuscate with custom options sqllexer -replace-digits=false -keep-json-path=true -input query.sql + +# Keep PostgreSQL IS predicate targets while replacing other booleans and NULLs +echo "SELECT true WHERE false IS NOT DISTINCT FROM true AND NULL IS NULL" | \ + sqllexer -dbms postgresql -keep-is-predicate ``` +`-keep-is-predicate` is disabled by default. For PostgreSQL, it preserves +`TRUE`, `FALSE`, and `NULL` on the right of `IS`, `IS NOT`, +`IS DISTINCT FROM`, and `IS NOT DISTINCT FROM`. Other literals still follow +`-replace-boolean` and `-replace-null`. The library option is +`WithKeepISPredicate(true)` and its JSON configuration key is `keep_is_predicate`. + +`pg_stat_statements` can parameterize boolean literals after `IS [NOT] DISTINCT +FROM`. For example, PostgreSQL 18.6 reports `SELECT $1 IS DISTINCT FROM $2` +for `SELECT TRUE IS DISTINCT FROM FALSE`. With this option enabled and positional +parameters replaced, the raw sample becomes `SELECT ? IS DISTINCT FROM FALSE` +while the representative text becomes `SELECT ? IS DISTINCT FROM ?`. These +queries therefore have different signatures. `IS TRUE` and `IS NULL` matched +in the same check. + ### Available Modes - **obfuscate** (default): Replace sensitive data with placeholders diff --git a/cmd/sqllexer/main.go b/cmd/sqllexer/main.go index 22e7b30..a723c58 100644 --- a/cmd/sqllexer/main.go +++ b/cmd/sqllexer/main.go @@ -17,6 +17,7 @@ type ObfuscatorConfig struct { ReplaceDigits bool ReplaceBoolean bool ReplaceNull bool + KeepISPredicate bool ReplaceBindParameter bool ReplacePositionalParameter bool DollarQuotedFunc bool @@ -53,6 +54,7 @@ func (c *ObfuscatorConfig) NewObfuscator() *sqllexer.Obfuscator { sqllexer.WithReplaceDigits(c.ReplaceDigits), sqllexer.WithReplaceBoolean(c.ReplaceBoolean), sqllexer.WithReplaceNull(c.ReplaceNull), + sqllexer.WithKeepISPredicate(c.KeepISPredicate), sqllexer.WithReplaceBindParameter(c.ReplaceBindParameter), sqllexer.WithReplacePositionalParameter(c.ReplacePositionalParameter), sqllexer.WithDollarQuotedFunc(c.DollarQuotedFunc), @@ -94,6 +96,7 @@ func parseFlags() *CLIConfig { flag.BoolVar(&cfg.Obfuscator.ReplaceDigits, "replace-digits", true, "Replace digits in identifiers with placeholders") flag.BoolVar(&cfg.Obfuscator.ReplaceBoolean, "replace-boolean", true, "Replace boolean values with placeholders") flag.BoolVar(&cfg.Obfuscator.ReplaceNull, "replace-null", true, "Replace NULL values with placeholders") + flag.BoolVar(&cfg.Obfuscator.KeepISPredicate, "keep-is-predicate", false, "Keep PostgreSQL TRUE, FALSE, and NULL after IS predicates") flag.BoolVar(&cfg.Obfuscator.ReplaceBindParameter, "replace-bind-parameter", false, "Replace bind parameters with placeholders") flag.BoolVar(&cfg.Obfuscator.ReplacePositionalParameter, "replace-positional-parameter", false, "Replace positional parameters ($1, $2, etc.) with placeholders") flag.BoolVar(&cfg.Obfuscator.DollarQuotedFunc, "dollar-quoted-func", false, "Obfuscate content inside $func$...$func$ blocks instead of replacing entirely") @@ -287,6 +290,8 @@ Obfuscator Flags: Replace boolean values with placeholders (default true) -replace-null Replace NULL values with placeholders (default true) + -keep-is-predicate + Keep PostgreSQL TRUE, FALSE, and NULL after IS predicates (default false) -replace-bind-parameter Replace bind parameters with placeholders (default false) -replace-positional-parameter diff --git a/dbms_test.go b/dbms_test.go index c66368d..1bc4705 100644 --- a/dbms_test.go +++ b/dbms_test.go @@ -114,6 +114,7 @@ func TestQueriesPerDBMS(t *testing.T) { WithReplacePositionalParameter(defaultObfuscatorConfig.ReplacePositionalParameter), WithReplaceBoolean(defaultObfuscatorConfig.ReplaceBoolean), WithReplaceNull(defaultObfuscatorConfig.ReplaceNull), + WithKeepISPredicate(defaultObfuscatorConfig.KeepISPredicate), WithKeepJsonPath(defaultObfuscatorConfig.KeepJsonPath), ) diff --git a/is_predicate.go b/is_predicate.go new file mode 100644 index 0000000..a83266a --- /dev/null +++ b/is_predicate.go @@ -0,0 +1,64 @@ +package sqllexer + +import "strings" + +type isPredicateState uint8 + +const ( + noISPredicate isPredicateState = iota + afterIS + afterISNot + afterISDistinct + afterISNotDistinct + afterISDistinctFrom +) + +// isPredicateContext follows only the keyword sequence of a PostgreSQL IS +// predicate. Whitespace and comments do not interrupt the sequence. +type isPredicateContext struct { + state isPredicateState +} + +func (c *isPredicateContext) keep(token *Token) bool { + if token.Type == SPACE || token.Type == COMMENT || token.Type == MULTILINE_COMMENT { + return false + } + + keyword := func(value string) bool { + return token.Type == KEYWORD && strings.EqualFold(token.Value, value) + } + operand := token.Type == BOOLEAN || token.Type == NULL + keep := false + switch c.state { + case afterIS: + switch { + case keyword("NOT"): + c.state = afterISNot + return false + case keyword("DISTINCT"): + c.state = afterISDistinct + return false + default: + keep = operand + } + case afterISNot: + if keyword("DISTINCT") { + c.state = afterISNotDistinct + return false + } + keep = operand + case afterISDistinct, afterISNotDistinct: + if keyword("FROM") { + c.state = afterISDistinctFrom + return false + } + case afterISDistinctFrom: + keep = operand + } + + c.state = noISPredicate + if keyword("IS") { + c.state = afterIS + } + return keep +} diff --git a/is_predicate_test.go b/is_predicate_test.go new file mode 100644 index 0000000..0bbdca7 --- /dev/null +++ b/is_predicate_test.go @@ -0,0 +1,179 @@ +package sqllexer + +import ( + "encoding/json" + "strconv" + "strings" + "testing" +) + +func checkISPredicateOutputs(t *testing.T, input string, obfuscator *Obfuscator, dbms DBMSType, wantObfuscated, wantNormalized string) { + t.Helper() + if got := obfuscator.Obfuscate(input, WithDBMS(dbms)); got != wantObfuscated { + t.Errorf("Obfuscate: got %q, want %q", got, wantObfuscated) + } + got, _, err := ObfuscateAndNormalize(input, obfuscator, NewNormalizer(), WithDBMS(dbms)) + if err != nil { + t.Fatal(err) + } + if got != wantNormalized { + t.Errorf("ObfuscateAndNormalize: got %q, want %q", got, wantNormalized) + } +} + +// The off goldens were captured from origin/main before this option was added. +func TestKeepISPredicateGoldens(t *testing.T) { + const input = "SELECT TRUE, NULL, a IS TRUE, b IS NOT FALSE, c IS NULL, d IS DISTINCT FROM TRUE, e IS NOT DISTINCT FROM NULL, TRUE = FALSE FROM t WHERE $1 IS TRUE" + tests := []struct { + name string + replaceBoolean bool + replaceNull bool + off string + on string + }{ + { + name: "neither", + off: "SELECT TRUE, NULL, a IS TRUE, b IS NOT FALSE, c IS NULL, d IS DISTINCT FROM TRUE, e IS NOT DISTINCT FROM NULL, TRUE = FALSE FROM t WHERE ? IS TRUE", + on: "SELECT TRUE, NULL, a IS TRUE, b IS NOT FALSE, c IS NULL, d IS DISTINCT FROM TRUE, e IS NOT DISTINCT FROM NULL, TRUE = FALSE FROM t WHERE ? IS TRUE", + }, + { + name: "null", + replaceNull: true, + off: "SELECT TRUE, ?, a IS TRUE, b IS NOT FALSE, c IS ?, d IS DISTINCT FROM TRUE, e IS NOT DISTINCT FROM ?, TRUE = FALSE FROM t WHERE ? IS TRUE", + on: "SELECT TRUE, ?, a IS TRUE, b IS NOT FALSE, c IS NULL, d IS DISTINCT FROM TRUE, e IS NOT DISTINCT FROM NULL, TRUE = FALSE FROM t WHERE ? IS TRUE", + }, + { + name: "boolean", + replaceBoolean: true, + off: "SELECT ?, NULL, a IS ?, b IS NOT ?, c IS NULL, d IS DISTINCT FROM ?, e IS NOT DISTINCT FROM NULL, ? = ? FROM t WHERE ? IS ?", + on: "SELECT ?, NULL, a IS TRUE, b IS NOT FALSE, c IS NULL, d IS DISTINCT FROM TRUE, e IS NOT DISTINCT FROM NULL, ? = ? FROM t WHERE ? IS TRUE", + }, + { + name: "both", + replaceBoolean: true, + replaceNull: true, + off: "SELECT ?, ?, a IS ?, b IS NOT ?, c IS ?, d IS DISTINCT FROM ?, e IS NOT DISTINCT FROM ?, ? = ? FROM t WHERE ? IS ?", + on: "SELECT ?, ?, a IS TRUE, b IS NOT FALSE, c IS NULL, d IS DISTINCT FROM TRUE, e IS NOT DISTINCT FROM NULL, ? = ? FROM t WHERE ? IS TRUE", + }, + } + + for _, tt := range tests { + for _, keep := range []bool{false, true} { + t.Run(tt.name+"/keep="+strconv.FormatBool(keep), func(t *testing.T) { + obfuscator := NewObfuscator( + WithReplaceBoolean(tt.replaceBoolean), + WithReplaceNull(tt.replaceNull), + WithReplacePositionalParameter(true), + WithKeepISPredicate(keep), + ) + want := tt.off + if keep { + want = tt.on + } + checkISPredicateOutputs(t, input, obfuscator, DBMSPostgres, want, want) + }) + } + } + + // Omitting the option must have exactly the same behavior as setting false. + checkISPredicateOutputs(t, input, NewObfuscator(WithReplaceBoolean(true), WithReplaceNull(true), WithReplacePositionalParameter(true)), DBMSPostgres, tests[3].off, tests[3].off) +} + +func TestKeepISPredicateForms(t *testing.T) { + for _, prefix := range []string{"IS", "IS NOT", "IS DISTINCT FROM", "IS NOT DISTINCT FROM"} { + for _, literal := range []string{"TRUE", "FALSE", "NULL"} { + input := "SELECT x " + prefix + " " + literal + t.Run(prefix+"/"+literal, func(t *testing.T) { + obfuscator := NewObfuscator(WithReplaceBoolean(true), WithReplaceNull(true), WithKeepISPredicate(true)) + checkISPredicateOutputs(t, input, obfuscator, DBMSPostgres, input, input) + }) + } + } +} + +func TestKeepISPredicateTokenBoundaries(t *testing.T) { + tests := []struct { + name string + input string + obfuscated string + normalized string + }{ + { + name: "casts aliases JSON and invalid bind", + input: "SELECT (x::boolean) IS /* a */ tRuE AS ok, NULL IS NOT FALSE, data ? 'IS TRUE' FROM t WHERE $1 IS TRUE AND x IS $1 AND x IS DISTINCT FROM 42 AND x IS NOT UNKNOWN", + obfuscated: "SELECT (x::boolean) IS /* a */ tRuE AS ok, ? IS NOT FALSE, data ? ? FROM t WHERE ? IS TRUE AND x IS ? AND x IS DISTINCT FROM ? AND x IS NOT UNKNOWN", + normalized: "SELECT ( x :: boolean ) IS tRuE, ? IS NOT FALSE, data ? ? FROM t WHERE ? IS TRUE AND x IS ? AND x IS DISTINCT FROM ? AND x IS NOT UNKNOWN", + }, + { + name: "comments strings and quoted identifiers", + input: `SELECT x IS /* NOT TRUE */ TRUE, x IS NOT /* IS TRUE */ DISTINCT /*hi*/ FROM FALSE, x IS DISTINCT FROM NULL, 'x IS TRUE', "IS TRUE" FROM t`, + obfuscated: `SELECT x IS /* NOT TRUE */ TRUE, x IS NOT /* IS TRUE */ DISTINCT /*hi*/ FROM FALSE, x IS DISTINCT FROM NULL, ?, "IS TRUE" FROM t`, + normalized: `SELECT x IS TRUE, x IS NOT DISTINCT FROM FALSE, x IS DISTINCT FROM NULL, ?, "IS TRUE" FROM t`, + }, + { + name: "binds and nonboolean distinct operand", + input: "SELECT $1 IS TRUE, x IS $1, x IS DISTINCT FROM $2, x IS DISTINCT FROM 42, x IS NOT UNKNOWN", + obfuscated: "SELECT ? IS TRUE, x IS ?, x IS DISTINCT FROM ?, x IS DISTINCT FROM ?, x IS NOT UNKNOWN", + normalized: "SELECT ? IS TRUE, x IS ?, x IS DISTINCT FROM ?, x IS DISTINCT FROM ?, x IS NOT UNKNOWN", + }, + { + name: "ordinary values", + input: "VALUES (TRUE, FALSE, NULL)", + obfuscated: "VALUES (?, ?, ?)", + normalized: "VALUES ( ? )", + }, + { + name: "nonboolean left operands", + input: "SELECT 42 IS NULL, 'value' IS NOT NULL", + obfuscated: "SELECT ? IS NULL, ? IS NOT NULL", + normalized: "SELECT ? IS NULL, ? IS NOT NULL", + }, + { + name: "case line comments and nested parentheses", + input: "SELECT (x iS -- comment IS NOT NULL\n tRuE) IS FALSE, y Is NoT DiStInCt -- comment FROM TRUE\n FrOm nUlL, TRUE::boolean IS FALSE", + obfuscated: "SELECT (x iS -- comment IS NOT NULL\n tRuE) IS FALSE, y Is NoT DiStInCt -- comment FROM TRUE\n FrOm nUlL, ?::boolean IS FALSE", + normalized: "SELECT ( x iS tRuE ) IS FALSE, y Is NoT DiStInCt FrOm nUlL, ? :: boolean IS FALSE", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + obfuscator := NewObfuscator(WithReplaceBoolean(true), WithReplaceNull(true), WithReplacePositionalParameter(true), WithKeepISPredicate(true)) + checkISPredicateOutputs(t, tt.input, obfuscator, DBMSPostgres, tt.obfuscated, tt.normalized) + }) + } +} + +// These representative texts were observed on PostgreSQL 18.6 with +// pg_stat_statements enabled. The mismatch is intentional for this option. +func TestKeepISPredicatePgStatStatementsRepresentative(t *testing.T) { + obfuscator := NewObfuscator(WithReplaceBoolean(true), WithReplaceNull(true), WithReplacePositionalParameter(true), WithKeepISPredicate(true)) + checkISPredicateOutputs(t, "SELECT TRUE IS DISTINCT FROM FALSE", obfuscator, DBMSPostgres, "SELECT ? IS DISTINCT FROM FALSE", "SELECT ? IS DISTINCT FROM FALSE") + checkISPredicateOutputs(t, "SELECT $1 IS DISTINCT FROM $2", obfuscator, DBMSPostgres, "SELECT ? IS DISTINCT FROM ?", "SELECT ? IS DISTINCT FROM ?") + checkISPredicateOutputs(t, "SELECT TRUE IS TRUE", obfuscator, DBMSPostgres, "SELECT ? IS TRUE", "SELECT ? IS TRUE") + checkISPredicateOutputs(t, "SELECT $1 IS TRUE", obfuscator, DBMSPostgres, "SELECT ? IS TRUE", "SELECT ? IS TRUE") +} + +func TestKeepISPredicateOnlyPostgres(t *testing.T) { + const input = "SELECT TRUE, x IS TRUE, x IS NOT DISTINCT FROM NULL" + const want = "SELECT ?, x IS ?, x IS NOT DISTINCT FROM ?" + for _, dbms := range []DBMSType{DBMSMySQL, DBMSOracle, DBMSSQLServer, DBMSSnowflake, ""} { + t.Run(string(dbms), func(t *testing.T) { + obfuscator := NewObfuscator(WithReplaceBoolean(true), WithReplaceNull(true), WithKeepISPredicate(true)) + checkISPredicateOutputs(t, input, obfuscator, dbms, want, want) + }) + } + checkISPredicateOutputs(t, input, NewObfuscator(WithReplaceBoolean(true), WithReplaceNull(true), WithKeepISPredicate(true)), DBMSPostgresAlias1, "SELECT ?, x IS TRUE, x IS NOT DISTINCT FROM NULL", "SELECT ?, x IS TRUE, x IS NOT DISTINCT FROM NULL") +} + +func TestKeepISPredicateConfigJSON(t *testing.T) { + got, err := json.Marshal(NewObfuscator(WithKeepISPredicate(true)).config) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(got), `"keep_is_predicate":true`) { + t.Fatalf("missing keep_is_predicate in %s", got) + } + if NewObfuscator().config.KeepISPredicate { + t.Fatal("KeepISPredicate must default to false") + } +} diff --git a/obfuscate_and_normalize.go b/obfuscate_and_normalize.go index 81857dd..c2a1b5c 100644 --- a/obfuscate_and_normalize.go +++ b/obfuscate_and_normalize.go @@ -4,8 +4,19 @@ package sqllexer // 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 + var predicate isPredicateContext + keepISPredicate := false + if obfuscator.config.KeepISPredicate { + lexerConfig := LexerConfig{} + for _, opt := range lexerOpts { + opt(&lexerConfig) + } + keepISPredicate = lexerConfig.DBMS == DBMSPostgres + } obfuscate := func(token *Token, lastValueToken *LastValueToken) { - obfuscator.ObfuscateTokenValue(token, lastValueToken, lexerOpts...) + if !keepISPredicate || !predicate.keep(token) { + obfuscator.ObfuscateTokenValue(token, lastValueToken, lexerOpts...) + } ec.maybeReplaceExtractField(token) ec.update(token) } diff --git a/obfuscator.go b/obfuscator.go index 4c00947..803ecd4 100644 --- a/obfuscator.go +++ b/obfuscator.go @@ -10,6 +10,7 @@ type obfuscatorConfig struct { ReplacePositionalParameter bool `json:"replace_positional_parameter"` ReplaceBoolean bool `json:"replace_boolean"` ReplaceNull bool `json:"replace_null"` + KeepISPredicate bool `json:"keep_is_predicate"` KeepJsonPath bool `json:"keep_json_path"` // by default, we replace json path with placeholder ReplaceBindParameter bool `json:"replace_bind_parameter"` } @@ -40,6 +41,14 @@ func WithReplaceNull(replaceNull bool) obfuscatorOption { } } +// WithKeepISPredicate preserves PostgreSQL TRUE, FALSE, and NULL predicate targets. +// The option is disabled by default and does not change other literal handling. +func WithKeepISPredicate(keepISPredicate bool) obfuscatorOption { + return func(c *obfuscatorConfig) { + c.KeepISPredicate = keepISPredicate + } +} + func WithDollarQuotedFunc(dollarQuotedFunc bool) obfuscatorOption { return func(c *obfuscatorConfig) { c.DollarQuotedFunc = dollarQuotedFunc @@ -92,13 +101,17 @@ func (o *Obfuscator) Obfuscate(input string, lexerOpts ...lexerOption) string { var lastValueToken *LastValueToken var ec extractContext + var predicate isPredicateContext + keepISPredicate := o.config.KeepISPredicate && lexer.config.DBMS == DBMSPostgres for { token := lexer.Scan() if token.Type == EOF { break } - o.ObfuscateTokenValue(token, lastValueToken, lexerOpts...) + if !keepISPredicate || !predicate.keep(token) { + o.ObfuscateTokenValue(token, lastValueToken, lexerOpts...) + } ec.maybeReplaceExtractField(token) obfuscatedSQL.WriteString(token.Value) From b4f8c2aba709134a57e7118b9e1986cb9338cfd9 Mon Sep 17 00:00:00 2001 From: Joel Marcotte Date: Wed, 30 Sep 2026 14:29:37 -0400 Subject: [PATCH 2/3] Defer DISTINCT FROM literal preservation --- README.md | 21 ++++++++++----------- cmd/sqllexer/main.go | 4 ++-- is_predicate.go | 23 ++--------------------- is_predicate_test.go | 37 ++++++++++++++++++++++++++----------- obfuscator.go | 2 +- 5 files changed, 41 insertions(+), 46 deletions(-) diff --git a/README.md b/README.md index a8f5f29..c7390f8 100644 --- a/README.md +++ b/README.md @@ -116,23 +116,22 @@ sqllexer -mode tokenize -input query.sql sqllexer -replace-digits=false -keep-json-path=true -input query.sql # Keep PostgreSQL IS predicate targets while replacing other booleans and NULLs -echo "SELECT true WHERE false IS NOT DISTINCT FROM true AND NULL IS NULL" | \ +echo "SELECT true WHERE false IS NOT true AND NULL IS NULL" | \ sqllexer -dbms postgresql -keep-is-predicate ``` `-keep-is-predicate` is disabled by default. For PostgreSQL, it preserves -`TRUE`, `FALSE`, and `NULL` on the right of `IS`, `IS NOT`, -`IS DISTINCT FROM`, and `IS NOT DISTINCT FROM`. Other literals still follow -`-replace-boolean` and `-replace-null`. The library option is +`TRUE`, `FALSE`, and `NULL` on the right of `IS` and `IS NOT`. Other literals, +including those after `IS [NOT] DISTINCT FROM`, still follow `-replace-boolean` +and `-replace-null`. The library option is `WithKeepISPredicate(true)` and its JSON configuration key is `keep_is_predicate`. -`pg_stat_statements` can parameterize boolean literals after `IS [NOT] DISTINCT -FROM`. For example, PostgreSQL 18.6 reports `SELECT $1 IS DISTINCT FROM $2` -for `SELECT TRUE IS DISTINCT FROM FALSE`. With this option enabled and positional -parameters replaced, the raw sample becomes `SELECT ? IS DISTINCT FROM FALSE` -while the representative text becomes `SELECT ? IS DISTINCT FROM ?`. These -queries therefore have different signatures. `IS TRUE` and `IS NULL` matched -in the same check. +`pg_stat_statements` can parameterize boolean literals after +`IS [NOT] DISTINCT FROM`. For example, PostgreSQL 18.6 reports +`SELECT $1 IS DISTINCT FROM $2` +for `SELECT TRUE IS DISTINCT FROM FALSE`. Preserving the right-hand `FALSE` +would give that raw sample a different signature from the representative text. +The option leaves both `DISTINCT FROM` forms on the existing replacement rules. ### Available Modes diff --git a/cmd/sqllexer/main.go b/cmd/sqllexer/main.go index a723c58..e01d0b2 100644 --- a/cmd/sqllexer/main.go +++ b/cmd/sqllexer/main.go @@ -96,7 +96,7 @@ func parseFlags() *CLIConfig { flag.BoolVar(&cfg.Obfuscator.ReplaceDigits, "replace-digits", true, "Replace digits in identifiers with placeholders") flag.BoolVar(&cfg.Obfuscator.ReplaceBoolean, "replace-boolean", true, "Replace boolean values with placeholders") flag.BoolVar(&cfg.Obfuscator.ReplaceNull, "replace-null", true, "Replace NULL values with placeholders") - flag.BoolVar(&cfg.Obfuscator.KeepISPredicate, "keep-is-predicate", false, "Keep PostgreSQL TRUE, FALSE, and NULL after IS predicates") + flag.BoolVar(&cfg.Obfuscator.KeepISPredicate, "keep-is-predicate", false, "Keep PostgreSQL TRUE, FALSE, and NULL after IS or IS NOT") flag.BoolVar(&cfg.Obfuscator.ReplaceBindParameter, "replace-bind-parameter", false, "Replace bind parameters with placeholders") flag.BoolVar(&cfg.Obfuscator.ReplacePositionalParameter, "replace-positional-parameter", false, "Replace positional parameters ($1, $2, etc.) with placeholders") flag.BoolVar(&cfg.Obfuscator.DollarQuotedFunc, "dollar-quoted-func", false, "Obfuscate content inside $func$...$func$ blocks instead of replacing entirely") @@ -291,7 +291,7 @@ Obfuscator Flags: -replace-null Replace NULL values with placeholders (default true) -keep-is-predicate - Keep PostgreSQL TRUE, FALSE, and NULL after IS predicates (default false) + Keep PostgreSQL TRUE, FALSE, and NULL after IS or IS NOT (default false) -replace-bind-parameter Replace bind parameters with placeholders (default false) -replace-positional-parameter diff --git a/is_predicate.go b/is_predicate.go index a83266a..99b624c 100644 --- a/is_predicate.go +++ b/is_predicate.go @@ -8,9 +8,6 @@ const ( noISPredicate isPredicateState = iota afterIS afterISNot - afterISDistinct - afterISNotDistinct - afterISDistinctFrom ) // isPredicateContext follows only the keyword sequence of a PostgreSQL IS @@ -31,28 +28,12 @@ func (c *isPredicateContext) keep(token *Token) bool { keep := false switch c.state { case afterIS: - switch { - case keyword("NOT"): + if keyword("NOT") { c.state = afterISNot return false - case keyword("DISTINCT"): - c.state = afterISDistinct - return false - default: - keep = operand - } - case afterISNot: - if keyword("DISTINCT") { - c.state = afterISNotDistinct - return false } keep = operand - case afterISDistinct, afterISNotDistinct: - if keyword("FROM") { - c.state = afterISDistinctFrom - return false - } - case afterISDistinctFrom: + case afterISNot: keep = operand } diff --git a/is_predicate_test.go b/is_predicate_test.go index 0bbdca7..96dd8cf 100644 --- a/is_predicate_test.go +++ b/is_predicate_test.go @@ -40,20 +40,20 @@ func TestKeepISPredicateGoldens(t *testing.T) { name: "null", replaceNull: true, off: "SELECT TRUE, ?, a IS TRUE, b IS NOT FALSE, c IS ?, d IS DISTINCT FROM TRUE, e IS NOT DISTINCT FROM ?, TRUE = FALSE FROM t WHERE ? IS TRUE", - on: "SELECT TRUE, ?, a IS TRUE, b IS NOT FALSE, c IS NULL, d IS DISTINCT FROM TRUE, e IS NOT DISTINCT FROM NULL, TRUE = FALSE FROM t WHERE ? IS TRUE", + on: "SELECT TRUE, ?, a IS TRUE, b IS NOT FALSE, c IS NULL, d IS DISTINCT FROM TRUE, e IS NOT DISTINCT FROM ?, TRUE = FALSE FROM t WHERE ? IS TRUE", }, { name: "boolean", replaceBoolean: true, off: "SELECT ?, NULL, a IS ?, b IS NOT ?, c IS NULL, d IS DISTINCT FROM ?, e IS NOT DISTINCT FROM NULL, ? = ? FROM t WHERE ? IS ?", - on: "SELECT ?, NULL, a IS TRUE, b IS NOT FALSE, c IS NULL, d IS DISTINCT FROM TRUE, e IS NOT DISTINCT FROM NULL, ? = ? FROM t WHERE ? IS TRUE", + on: "SELECT ?, NULL, a IS TRUE, b IS NOT FALSE, c IS NULL, d IS DISTINCT FROM ?, e IS NOT DISTINCT FROM NULL, ? = ? FROM t WHERE ? IS TRUE", }, { name: "both", replaceBoolean: true, replaceNull: true, off: "SELECT ?, ?, a IS ?, b IS NOT ?, c IS ?, d IS DISTINCT FROM ?, e IS NOT DISTINCT FROM ?, ? = ? FROM t WHERE ? IS ?", - on: "SELECT ?, ?, a IS TRUE, b IS NOT FALSE, c IS NULL, d IS DISTINCT FROM TRUE, e IS NOT DISTINCT FROM NULL, ? = ? FROM t WHERE ? IS TRUE", + on: "SELECT ?, ?, a IS TRUE, b IS NOT FALSE, c IS NULL, d IS DISTINCT FROM ?, e IS NOT DISTINCT FROM ?, ? = ? FROM t WHERE ? IS TRUE", }, } @@ -80,7 +80,7 @@ func TestKeepISPredicateGoldens(t *testing.T) { } func TestKeepISPredicateForms(t *testing.T) { - for _, prefix := range []string{"IS", "IS NOT", "IS DISTINCT FROM", "IS NOT DISTINCT FROM"} { + for _, prefix := range []string{"IS", "IS NOT"} { for _, literal := range []string{"TRUE", "FALSE", "NULL"} { input := "SELECT x " + prefix + " " + literal t.Run(prefix+"/"+literal, func(t *testing.T) { @@ -91,6 +91,19 @@ func TestKeepISPredicateForms(t *testing.T) { } } +func TestKeepISPredicateDoesNotChangeDistinctFrom(t *testing.T) { + for _, prefix := range []string{"IS DISTINCT FROM", "IS NOT DISTINCT FROM"} { + for _, literal := range []string{"TRUE", "FALSE", "NULL"} { + input := "SELECT x " + prefix + " " + literal + want := "SELECT x " + prefix + " ?" + t.Run(prefix+"/"+literal, func(t *testing.T) { + obfuscator := NewObfuscator(WithReplaceBoolean(true), WithReplaceNull(true), WithKeepISPredicate(true)) + checkISPredicateOutputs(t, input, obfuscator, DBMSPostgres, want, want) + }) + } + } +} + func TestKeepISPredicateTokenBoundaries(t *testing.T) { tests := []struct { name string @@ -107,8 +120,8 @@ func TestKeepISPredicateTokenBoundaries(t *testing.T) { { name: "comments strings and quoted identifiers", input: `SELECT x IS /* NOT TRUE */ TRUE, x IS NOT /* IS TRUE */ DISTINCT /*hi*/ FROM FALSE, x IS DISTINCT FROM NULL, 'x IS TRUE', "IS TRUE" FROM t`, - obfuscated: `SELECT x IS /* NOT TRUE */ TRUE, x IS NOT /* IS TRUE */ DISTINCT /*hi*/ FROM FALSE, x IS DISTINCT FROM NULL, ?, "IS TRUE" FROM t`, - normalized: `SELECT x IS TRUE, x IS NOT DISTINCT FROM FALSE, x IS DISTINCT FROM NULL, ?, "IS TRUE" FROM t`, + obfuscated: `SELECT x IS /* NOT TRUE */ TRUE, x IS NOT /* IS TRUE */ DISTINCT /*hi*/ FROM ?, x IS DISTINCT FROM ?, ?, "IS TRUE" FROM t`, + normalized: `SELECT x IS TRUE, x IS NOT DISTINCT FROM ?, x IS DISTINCT FROM ?, ?, "IS TRUE" FROM t`, }, { name: "binds and nonboolean distinct operand", @@ -131,8 +144,8 @@ func TestKeepISPredicateTokenBoundaries(t *testing.T) { { name: "case line comments and nested parentheses", input: "SELECT (x iS -- comment IS NOT NULL\n tRuE) IS FALSE, y Is NoT DiStInCt -- comment FROM TRUE\n FrOm nUlL, TRUE::boolean IS FALSE", - obfuscated: "SELECT (x iS -- comment IS NOT NULL\n tRuE) IS FALSE, y Is NoT DiStInCt -- comment FROM TRUE\n FrOm nUlL, ?::boolean IS FALSE", - normalized: "SELECT ( x iS tRuE ) IS FALSE, y Is NoT DiStInCt FrOm nUlL, ? :: boolean IS FALSE", + obfuscated: "SELECT (x iS -- comment IS NOT NULL\n tRuE) IS FALSE, y Is NoT DiStInCt -- comment FROM TRUE\n FrOm ?, ?::boolean IS FALSE", + normalized: "SELECT ( x iS tRuE ) IS FALSE, y Is NoT DiStInCt FrOm ?, ? :: boolean IS FALSE", }, } for _, tt := range tests { @@ -144,11 +157,13 @@ func TestKeepISPredicateTokenBoundaries(t *testing.T) { } // These representative texts were observed on PostgreSQL 18.6 with -// pg_stat_statements enabled. The mismatch is intentional for this option. +// pg_stat_statements enabled. Excluding DISTINCT FROM keeps these pairs aligned. func TestKeepISPredicatePgStatStatementsRepresentative(t *testing.T) { obfuscator := NewObfuscator(WithReplaceBoolean(true), WithReplaceNull(true), WithReplacePositionalParameter(true), WithKeepISPredicate(true)) - checkISPredicateOutputs(t, "SELECT TRUE IS DISTINCT FROM FALSE", obfuscator, DBMSPostgres, "SELECT ? IS DISTINCT FROM FALSE", "SELECT ? IS DISTINCT FROM FALSE") + checkISPredicateOutputs(t, "SELECT TRUE IS DISTINCT FROM FALSE", obfuscator, DBMSPostgres, "SELECT ? IS DISTINCT FROM ?", "SELECT ? IS DISTINCT FROM ?") checkISPredicateOutputs(t, "SELECT $1 IS DISTINCT FROM $2", obfuscator, DBMSPostgres, "SELECT ? IS DISTINCT FROM ?", "SELECT ? IS DISTINCT FROM ?") + checkISPredicateOutputs(t, "SELECT TRUE IS NOT DISTINCT FROM FALSE", obfuscator, DBMSPostgres, "SELECT ? IS NOT DISTINCT FROM ?", "SELECT ? IS NOT DISTINCT FROM ?") + checkISPredicateOutputs(t, "SELECT $1 IS NOT DISTINCT FROM $2", obfuscator, DBMSPostgres, "SELECT ? IS NOT DISTINCT FROM ?", "SELECT ? IS NOT DISTINCT FROM ?") checkISPredicateOutputs(t, "SELECT TRUE IS TRUE", obfuscator, DBMSPostgres, "SELECT ? IS TRUE", "SELECT ? IS TRUE") checkISPredicateOutputs(t, "SELECT $1 IS TRUE", obfuscator, DBMSPostgres, "SELECT ? IS TRUE", "SELECT ? IS TRUE") } @@ -162,7 +177,7 @@ func TestKeepISPredicateOnlyPostgres(t *testing.T) { checkISPredicateOutputs(t, input, obfuscator, dbms, want, want) }) } - checkISPredicateOutputs(t, input, NewObfuscator(WithReplaceBoolean(true), WithReplaceNull(true), WithKeepISPredicate(true)), DBMSPostgresAlias1, "SELECT ?, x IS TRUE, x IS NOT DISTINCT FROM NULL", "SELECT ?, x IS TRUE, x IS NOT DISTINCT FROM NULL") + checkISPredicateOutputs(t, input, NewObfuscator(WithReplaceBoolean(true), WithReplaceNull(true), WithKeepISPredicate(true)), DBMSPostgresAlias1, "SELECT ?, x IS TRUE, x IS NOT DISTINCT FROM ?", "SELECT ?, x IS TRUE, x IS NOT DISTINCT FROM ?") } func TestKeepISPredicateConfigJSON(t *testing.T) { diff --git a/obfuscator.go b/obfuscator.go index 803ecd4..9c856a3 100644 --- a/obfuscator.go +++ b/obfuscator.go @@ -41,7 +41,7 @@ func WithReplaceNull(replaceNull bool) obfuscatorOption { } } -// WithKeepISPredicate preserves PostgreSQL TRUE, FALSE, and NULL predicate targets. +// WithKeepISPredicate preserves PostgreSQL TRUE, FALSE, and NULL after IS or IS NOT. // The option is disabled by default and does not change other literal handling. func WithKeepISPredicate(keepISPredicate bool) obfuscatorOption { return func(c *obfuscatorConfig) { From 7d8054a7e36a6d7bcecd8d4d45e5adc8f44065f1 Mon Sep 17 00:00:00 2001 From: Joel Marcotte Date: Wed, 30 Sep 2026 14:40:29 -0400 Subject: [PATCH 3/3] Avoid allocation in IS predicate normalization path --- is_predicate.go | 35 ++++++++++++++----------------- is_predicate_bench_test.go | 42 ++++++++++++++++++++++++++++++++++++++ normalizer.go | 10 ++++----- obfuscate_and_normalize.go | 12 +++-------- 4 files changed, 65 insertions(+), 34 deletions(-) create mode 100644 is_predicate_bench_test.go diff --git a/is_predicate.go b/is_predicate.go index 99b624c..d8d4b50 100644 --- a/is_predicate.go +++ b/is_predicate.go @@ -1,7 +1,5 @@ package sqllexer -import "strings" - type isPredicateState uint8 const ( @@ -17,29 +15,26 @@ type isPredicateContext struct { } func (c *isPredicateContext) keep(token *Token) bool { - if token.Type == SPACE || token.Type == COMMENT || token.Type == MULTILINE_COMMENT { + switch token.Type { + case SPACE, COMMENT, MULTILINE_COMMENT: return false - } - - keyword := func(value string) bool { - return token.Type == KEYWORD && strings.EqualFold(token.Value, value) - } - operand := token.Type == BOOLEAN || token.Type == NULL - keep := false - switch c.state { - case afterIS: - if keyword("NOT") { + case KEYWORD: + // The lexer classifies SQL keywords from ASCII letters. + value := token.Value + if len(value) == 2 && value[0]|0x20 == 'i' && value[1]|0x20 == 's' { + c.state = afterIS + return false + } + if c.state == afterIS && len(value) == 3 && value[0]|0x20 == 'n' && value[1]|0x20 == 'o' && value[2]|0x20 == 't' { c.state = afterISNot return false } - keep = operand - case afterISNot: - keep = operand + case BOOLEAN, NULL: + keep := c.state == afterIS || c.state == afterISNot + c.state = noISPredicate + return keep } c.state = noISPredicate - if keyword("IS") { - c.state = afterIS - } - return keep + return false } diff --git a/is_predicate_bench_test.go b/is_predicate_bench_test.go new file mode 100644 index 0000000..87b9c39 --- /dev/null +++ b/is_predicate_bench_test.go @@ -0,0 +1,42 @@ +package sqllexer + +import "testing" + +var isPredicateBenchResult string + +func BenchmarkKeepISPredicate(b *testing.B) { + queries := []struct { + name string + sql string + }{ + {"no_predicate", "SELECT id, flag FROM events WHERE id = 42 AND status = 'ready'"}, + {"one_predicate", "SELECT id FROM events WHERE enabled IS TRUE"}, + {"mixed", "SELECT TRUE, e.enabled IS /* signal */ TRUE, e.deleted_at IS NOT NULL, e.payload ? 'flag' FROM events e WHERE e.id = $1 AND e.flag IS DISTINCT FROM FALSE"}, + } + lexerOpt := WithDBMS(DBMSPostgres) + normalizer := NewNormalizer() + for _, query := range queries { + for _, keep := range []struct { + name string + value bool + }{{"off", false}, {"on", true}} { + obfuscator := NewObfuscator(WithReplaceBoolean(true), WithReplaceNull(true), WithKeepISPredicate(keep.value)) + b.Run("Obfuscate/"+query.name+"/"+keep.name, func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + isPredicateBenchResult = obfuscator.Obfuscate(query.sql, lexerOpt) + } + }) + b.Run("ObfuscateAndNormalize/"+query.name+"/"+keep.name, func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + var err error + isPredicateBenchResult, _, err = ObfuscateAndNormalize(query.sql, obfuscator, normalizer, lexerOpt) + if err != nil { + b.Fatal(err) + } + } + }) + } + } +} diff --git a/normalizer.go b/normalizer.go index a893c43..bb8718f 100644 --- a/normalizer.go +++ b/normalizer.go @@ -196,13 +196,13 @@ func (n *Normalizer) normalizeToken(lexer *Lexer, normalizedSQLBuilder *strings. } func (n *Normalizer) Normalize(input string, lexerOpts ...lexerOption) (normalizedSQL string, statementMetadata *StatementMetadata, err error) { - return n.normalize(input, nil, lexerOpts...) + return n.normalize(New(input, lexerOpts...), nil, lexerOpts...) } -// normalize is the internal implementation that handles the common normalization logic. -// preProcessToken is an optional function to process tokens before normalization (e.g., obfuscation). -func (n *Normalizer) normalize(input string, preProcessToken func(*Token, *LastValueToken), lexerOpts ...lexerOption) (normalizedSQL string, statementMetadata *StatementMetadata, err error) { - lexer := New(input, lexerOpts...) +// normalize processes a prepared lexer so the combined path can use its DBMS configuration. +// preProcessToken optionally handles tokens before normalization (e.g., obfuscation). +func (n *Normalizer) normalize(lexer *Lexer, preProcessToken func(*Token, *LastValueToken), lexerOpts ...lexerOption) (normalizedSQL string, statementMetadata *StatementMetadata, err error) { + input := lexer.src var normalizedSQLBuilder strings.Builder normalizedSQLBuilder.Grow(len(input)) diff --git a/obfuscate_and_normalize.go b/obfuscate_and_normalize.go index c2a1b5c..13687e1 100644 --- a/obfuscate_and_normalize.go +++ b/obfuscate_and_normalize.go @@ -5,14 +5,8 @@ package sqllexer func ObfuscateAndNormalize(input string, obfuscator *Obfuscator, normalizer *Normalizer, lexerOpts ...lexerOption) (normalizedSQL string, statementMetadata *StatementMetadata, err error) { var ec extractContext var predicate isPredicateContext - keepISPredicate := false - if obfuscator.config.KeepISPredicate { - lexerConfig := LexerConfig{} - for _, opt := range lexerOpts { - opt(&lexerConfig) - } - keepISPredicate = lexerConfig.DBMS == DBMSPostgres - } + lexer := New(input, lexerOpts...) + keepISPredicate := obfuscator.config.KeepISPredicate && lexer.config.DBMS == DBMSPostgres obfuscate := func(token *Token, lastValueToken *LastValueToken) { if !keepISPredicate || !predicate.keep(token) { obfuscator.ObfuscateTokenValue(token, lastValueToken, lexerOpts...) @@ -20,5 +14,5 @@ func ObfuscateAndNormalize(input string, obfuscator *Obfuscator, normalizer *Nor ec.maybeReplaceExtractField(token) ec.update(token) } - return normalizer.normalize(input, obfuscate, lexerOpts...) + return normalizer.normalize(lexer, obfuscate, lexerOpts...) }