Add opt-in PostgreSQL IS predicate literal preservation

README.md

@@ -114,8 +114,25 @@ 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 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` 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`. 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
 
 - **obfuscate** (default): Replace sensitive data with placeholders

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 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")
@@ -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 or IS NOT (default false)
   -replace-bind-parameter
         Replace bind parameters with placeholders (default false)
   -replace-positional-parameter

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

is_predicate.go

@@ -0,0 +1,40 @@
+package sqllexer
+
+type isPredicateState uint8
+
+const (
+	noISPredicate isPredicateState = iota
+	afterIS
+	afterISNot
+)
+
+// 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 {
+	switch token.Type {
+	case SPACE, COMMENT, MULTILINE_COMMENT:
+		return false
+	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
+		}
+	case BOOLEAN, NULL:
+		keep := c.state == afterIS || c.state == afterISNot
+		c.state = noISPredicate
+		return keep
+	}
+
+	c.state = noISPredicate
+	return false
+}

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

is_predicate_test.go

@@ -0,0 +1,194 @@
+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 ?, 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 ?, 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 ?, e IS NOT DISTINCT FROM ?, ? = ? 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"} {
+		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 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
+		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 ?, 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",
+			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 ?, ?::boolean IS FALSE",
+			normalized: "SELECT ( x iS tRuE ) IS FALSE, y Is NoT DiStInCt FrOm ?, ? :: 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. 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 ?", "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")
+}
+
+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 ?", "SELECT ?, x IS TRUE, x IS NOT DISTINCT FROM ?")
+}
+
+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")
+	}
+}

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

obfuscate_and_normalize.go

@@ -4,10 +4,15 @@ 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
+	lexer := New(input, lexerOpts...)
+	keepISPredicate := obfuscator.config.KeepISPredicate && lexer.config.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)
 	}
-	return normalizer.normalize(input, obfuscate, lexerOpts...)
+	return normalizer.normalize(lexer, obfuscate, lexerOpts...)
 }