From c37c31fd685a080aba3f077dc83bdba0fd4b4103 Mon Sep 17 00:00:00 2001
From: Jason Petersen <jason.petersen@datadoghq.com>
Date: Mon, 5 Oct 2026 10:42:23 -0600
Subject: [PATCH] Handle Db2 LUW string literals without backslash escapes

---
 README.md            |  1 +
 cmd/sqllexer/main.go |  2 +-
 db2_test.go          | 56 ++++++++++++++++++++++++++++++++++++++++++++
 sqllexer.go          |  4 ++--
 sqllexer_utils.go    |  2 ++
 5 files changed, 62 insertions(+), 3 deletions(-)
 create mode 100644 db2_test.go

diff --git a/README.md b/README.md
index a9e0c24..b893727 100644
--- a/README.md
+++ b/README.md
@@ -130,6 +130,7 @@ Use the `-dbms` flag to specify the database type:
 - `mysql` - MySQL
 - `oracle` - Oracle
 - `snowflake` - Snowflake
+- `ibm_db2` - IBM Db2 LUW
 
 ## Testing
 
diff --git a/cmd/sqllexer/main.go b/cmd/sqllexer/main.go
index 22e7b30..35e37d7 100644
--- a/cmd/sqllexer/main.go
+++ b/cmd/sqllexer/main.go
@@ -87,7 +87,7 @@ func parseFlags() *CLIConfig {
 	flag.StringVar(&cfg.Mode, "mode", "obfuscate_and_normalize", "Operation mode: obfuscate, normalize, tokenize, obfuscate_and_normalize")
 	flag.StringVar(&cfg.InputFile, "input", "", "Input file (default: stdin)")
 	flag.StringVar(&cfg.OutputFile, "output", "", "Output file (default: stdout)")
-	flag.StringVar(&cfg.DBMS, "dbms", "", "Database type: mssql, postgresql, mysql, oracle, snowflake")
+	flag.StringVar(&cfg.DBMS, "dbms", "", "Database type: mssql, postgresql, mysql, oracle, snowflake, ibm_db2")
 	flag.BoolVar(&cfg.WithMetadata, "with-metadata", false, "Output result with metadata as JSON (normalize and obfuscate_and_normalize modes)")
 
 	// Obfuscator options
diff --git a/db2_test.go b/db2_test.go
new file mode 100644
index 0000000..178aa17
--- /dev/null
+++ b/db2_test.go
@@ -0,0 +1,56 @@
+package sqllexer
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// Db2 ordinary strings use doubled quotes and treat backslashes literally.
+// Check boundaries so a closing quote cannot swallow a comma or expose the next value.
+func TestDb2StringBoundaries(t *testing.T) {
+	for _, tc := range []struct {
+		input string
+		want  []TokenSpec
+	}{
+		{`'path\','probe_alpha'`, []TokenSpec{{STRING, `'path\'`}, {PUNCTUATION, ","}, {STRING, `'probe_alpha'`}}},
+		{`'probe_alpha''suffix'`, []TokenSpec{{STRING, `'probe_alpha''suffix'`}}},
+		{`'path\''秘密',17`, []TokenSpec{{STRING, `'path\''秘密'`}, {PUNCTUATION, ","}, {NUMBER, "17"}}},
+		{`''||''''`, []TokenSpec{{STRING, `''`}, {OPERATOR, "||"}, {STRING, `''''`}}},
+		{`'path\\',?`, []TokenSpec{{STRING, `'path\\'`}, {PUNCTUATION, ","}, {OPERATOR, "?"}}},
+		{`'unfinished\`, []TokenSpec{{INCOMPLETE_STRING, `'unfinished\`}}},
+	} {
+		t.Run(tc.input, func(t *testing.T) {
+			lexer := New(tc.input, WithDBMS(DBMSType("ibm_db2")))
+			for _, want := range tc.want {
+				got := lexer.Scan()
+				require.Equal(t, want.Type, got.Type)
+				require.Equal(t, want.Value, got.Value)
+			}
+			require.Equal(t, EOF, lexer.Scan().Type)
+		})
+	}
+}
+
+// Literal contents must neither survive obfuscation nor change the normalized
+// text from which the Agent's callers compute query signatures.
+func TestDb2LiteralNormalization(t *testing.T) {
+	for _, tc := range []struct {
+		first, second, want string
+	}{
+		{`VALUES 'path\', 'probe_alpha'`, `VALUES 'other', 'probe_beta'`, `VALUES ?, ?`},
+		{`VALUES 'probe_alpha''suffix'`, `VALUES 'probe_beta'`, `VALUES ?`},
+		{`VALUES 'path\''秘密', 'probe_alpha'`, `VALUES 'plain', 'probe_beta'`, `VALUES ?, ?`},
+		{`VALUES ''`, `VALUES ''''`, `VALUES ?`},
+		{`VALUES 'path\\'`, `VALUES 'plain'`, `VALUES ?`},
+		{`VALUES 'a' || 'b'`, `VALUES 'c' || 'd'`, `VALUES ? || ?`},
+	} {
+		t.Run(tc.first, func(t *testing.T) {
+			for _, query := range []string{tc.first, tc.second} {
+				got, _, err := ObfuscateAndNormalize(query, NewObfuscator(), NewNormalizer(), WithDBMS(DBMSType("ibm_db2")))
+				require.NoError(t, err)
+				require.Equal(t, tc.want, got)
+			}
+		})
+	}
+}
diff --git a/sqllexer.go b/sqllexer.go
index 630c3a0..a680d4c 100644
--- a/sqllexer.go
+++ b/sqllexer.go
@@ -356,12 +356,12 @@ func (s *Lexer) scanStringWithDelimiter(delimiter rune) *Token {
 // stringBackslashEscapes reports whether a backslash inside the literal that
 // starts at s.start escapes the following character.
 //
-// SQL Server and Oracle never treat backslash as an escape. PostgreSQL does
+// SQL Server, Oracle, and Db2 never treat backslash as an escape. PostgreSQL does
 // so only for escape strings (E'...' / e'...'). Ordinary PostgreSQL strings
 // follow standard_conforming_strings: a quote is escaped by doubling it.
 func (s *Lexer) stringBackslashEscapes() bool {
 	switch s.config.DBMS {
-	case DBMSSQLServer, DBMSOracle:
+	case DBMSSQLServer, DBMSOracle, DBMSDB2:
 		return false
 	case DBMSPostgres:
 		return s.hasPostgresEscapeStringPrefix()
diff --git a/sqllexer_utils.go b/sqllexer_utils.go
index bdbca25..74e95aa 100644
--- a/sqllexer_utils.go
+++ b/sqllexer_utils.go
@@ -19,6 +19,8 @@ const (
 	DBMSMySQL DBMSType = "mysql"
 	// DBMSOracle is a Oracle Server
 	DBMSOracle DBMSType = "oracle"
+	// DBMSDB2 is IBM Db2 LUW. The name matches the Agent integration.
+	DBMSDB2 DBMSType = "ibm_db2"
 	// DBMSSnowflake is a Snowflake Server
 	DBMSSnowflake DBMSType = "snowflake"
 )