diff --git a/README.md b/README.md index a9e0c24..c7390f8 100644 --- a/README.md +++ b/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 diff --git a/cmd/sqllexer/main.go b/cmd/sqllexer/main.go index 22e7b30..e01d0b2 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 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 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..d8d4b50 --- /dev/null +++ b/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 +} 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/is_predicate_test.go b/is_predicate_test.go new file mode 100644 index 0000000..96dd8cf --- /dev/null +++ b/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") + } +} 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 81857dd..13687e1 100644 --- a/obfuscate_and_normalize.go +++ b/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...) } diff --git a/obfuscator.go b/obfuscator.go index 4c00947..9c856a3 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 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) { + 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)