reduce unnecessary calls to trimQuotes when KeepIdentifierQuotation is set

normalizer.go

@@ -165,29 +165,31 @@ func (n *Normalizer) collectMetadata(token *Token, lastToken *Token, statementMe
 		// Collect comments
 		statementMetadata.Comments = append(statementMetadata.Comments, token.Value)
 	} else if token.Type == IDENT || token.Type == QUOTED_IDENT || token.Type == FUNCTION {
-		tokenVal := token.Value
 		if token.Type == QUOTED_IDENT {
-			// We always want to trim the quotes for collected metadata such as table names
-			// This is because the metadata is used as tags, and we don't want them to be normalized as underscores later on
-			tokenVal = trimQuotes(tokenVal, tokenVal[0:1], tokenVal[len(tokenVal)-1:])
 			if !n.config.KeepIdentifierQuotation {
-				token.Value = tokenVal
+				token.Value = trimQuotes(token.Value, token.Value[0:1], token.Value[len(token.Value)-1:])
 			}
 		}
-		if n.config.CollectCommands && isCommand(tokenVal) {
+		if n.config.CollectCommands && isCommand(token.Value) {
 			// Collect commands
-			statementMetadata.Commands = append(statementMetadata.Commands, strings.ToUpper(tokenVal))
+			statementMetadata.Commands = append(statementMetadata.Commands, strings.ToUpper(token.Value))
 		} else if isWith(lastToken.Value) && token.Type == IDENT {
 			// Collect CTEs so we can skip them later in table collection
-			ctes[tokenVal] = true
-		} else if n.config.CollectTables && isTableIndicator(lastToken.Value) && !isSQLKeyword(tokenVal) {
+			ctes[token.Value] = true
+		} else if n.config.CollectTables && isTableIndicator(lastToken.Value) && !isSQLKeyword(token.Value) {
 			// Collect table names the token is not a CTE
-			if _, ok := ctes[tokenVal]; !ok {
-				statementMetadata.Tables = append(statementMetadata.Tables, tokenVal)
+			if _, ok := ctes[token.Value]; !ok {
+				table := token.Value
+				// Remove quotes from table name if KeepIdentifierQuotation is false
+				// Quotes need to be removed from the table name because the table names are used as tags
+				if token.Type == QUOTED_IDENT && n.config.KeepIdentifierQuotation {
+					table = trimQuotes(token.Value, token.Value[0:1], token.Value[len(token.Value)-1:])
+				}
+				statementMetadata.Tables = append(statementMetadata.Tables, table)
 			}
 		} else if n.config.CollectProcedure && isProcedure(lastToken) {
 			// Collect procedure names
-			statementMetadata.Procedures = append(statementMetadata.Procedures, tokenVal)
+			statementMetadata.Procedures = append(statementMetadata.Procedures, token.Value)
 		}
 	}
 }

normalizer_test.go

@@ -829,7 +829,7 @@ func TestNormalizeDeobfuscatedSQL(t *testing.T) {
 			input:    "SELECT * FROM `public`.`users` WHERE id = ?",
 			expected: "SELECT * FROM `public`.`users` WHERE id = ?",
 			statementMetadata: StatementMetadata{
-				Tables:     []string{`public.users`},
+				Tables:     []string{"public.users"},
 				Comments:   []string{},
 				Commands:   []string{"SELECT"},
 				Procedures: []string{},
@@ -887,6 +887,27 @@ func TestNormalizeDeobfuscatedSQL(t *testing.T) {
 				WithDBMS(DBMSSQLServer),
 			},
 		},
+		{
+			input:    `SELECT * FROM [public].[my users] WHERE id = ?`,
+			expected: `SELECT * FROM [public].[my users] WHERE id = ?`,
+			statementMetadata: StatementMetadata{
+				Tables:     []string{`public.my users`},
+				Comments:   []string{},
+				Commands:   []string{"SELECT"},
+				Procedures: []string{},
+				Size:       21,
+			},
+			normalizationConfig: &normalizerConfig{
+				CollectComments:         true,
+				CollectCommands:         true,
+				CollectTables:           true,
+				KeepSQLAlias:            true,
+				KeepIdentifierQuotation: true,
+			},
+			lexerOptions: []lexerOption{
+				WithDBMS(DBMSSQLServer),
+			},
+		},
 	}
 
 	for _, test := range tests {

obfuscate_and_normalize_bench_test.go

@@ -118,16 +118,17 @@ ORDER BY
 	var backtickQuotedQuery = "SELECT `orders`.`OrderID`, `customers`.`CustomerName`, `products`.`ProductName`, `order_details`.`Quantity`, `order_details`.`UnitPrice`, (`order_details`.`Quantity` * `order_details`.`UnitPrice`) AS `TotalPrice`, `orders`.`OrderDate`, `orders`.`ShippedDate`, CASE WHEN `orders`.`ShippedDate` IS NULL THEN 'Pending' ELSE 'Shipped' END AS `OrderStatus` FROM `orders` INNER JOIN `customers` ON `orders`.`CustomerID` = `customers`.`CustomerID` INNER JOIN `order_details` ON `orders`.`OrderID` = `order_details`.`OrderID` INNER JOIN `products` ON `order_details`.`ProductID` = `products`.`ProductID` WHERE `orders`.`OrderDate` >= '2024-01-01' AND `orders`.`OrderDate` <= '2024-12-31' AND `customers`.`Region` = 'North America' GROUP BY `orders`.`OrderID`, `customers`.`CustomerName`, `products`.`ProductName`, `order_details`.`Quantity`, `order_details`.`UnitPrice`, `orders`.`OrderDate`, `orders`.`ShippedDate` HAVING SUM(`order_details`.`Quantity`) > 10 ORDER BY `orders`.`OrderDate` DESC;"
 
 	benchmarks := []struct {
-		name  string
-		query string
+		name         string
+		query        string
+		lexerOptions []lexerOption
 	}{
-		{"Escaping", `INSERT INTO delayed_jobs (attempts, created_at, failed_at, handler, last_error, locked_at, locked_by, priority, queue, run_at, updated_at) VALUES (0, '2016-12-04 17:09:59', NULL, '--- !ruby/object:Delayed::PerformableMethod\nobject: !ruby/object:Item\n  store:\n  - a simple string\n  - an \'escaped \' string\n  - another \'escaped\' string\n  - 42\n  string: a string with many \\\\\'escapes\\\\\'\nmethod_name: :show_store\nargs: []\n', NULL, NULL, NULL, 0, NULL, '2016-12-04 17:09:59', '2016-12-04 17:09:59')`},
-		{"Grouping", `INSERT INTO delayed_jobs (created_at, failed_at, handler) VALUES (0, '2016-12-04 17:09:59', NULL), (0, '2016-12-04 17:09:59', NULL), (0, '2016-12-04 17:09:59', NULL), (0, '2016-12-04 17:09:59', NULL)`},
-		{"Large", LargeQuery},
-		{"Complex", ComplexQuery},
-		{"SuperLarge", fmt.Sprintf(superLargeQuery, 1)},
-		{"BracketQuoted", bracketQuotedQuery},
-		{"BacktickQuoted", backtickQuotedQuery},
+		{"Escaping", `INSERT INTO delayed_jobs (attempts, created_at, failed_at, handler, last_error, locked_at, locked_by, priority, queue, run_at, updated_at) VALUES (0, '2016-12-04 17:09:59', NULL, '--- !ruby/object:Delayed::PerformableMethod\nobject: !ruby/object:Item\n  store:\n  - a simple string\n  - an \'escaped \' string\n  - another \'escaped\' string\n  - 42\n  string: a string with many \\\\\'escapes\\\\\'\nmethod_name: :show_store\nargs: []\n', NULL, NULL, NULL, 0, NULL, '2016-12-04 17:09:59', '2016-12-04 17:09:59')`, nil},
+		{"Grouping", `INSERT INTO delayed_jobs (created_at, failed_at, handler) VALUES (0, '2016-12-04 17:09:59', NULL), (0, '2016-12-04 17:09:59', NULL), (0, '2016-12-04 17:09:59', NULL), (0, '2016-12-04 17:09:59', NULL)`, nil},
+		{"Large", LargeQuery, nil},
+		{"Complex", ComplexQuery, nil},
+		{"SuperLarge", fmt.Sprintf(superLargeQuery, 1), nil},
+		{"BracketQuoted", bracketQuotedQuery, []lexerOption{WithDBMS(DBMSSQLServer)}},
+		{"BacktickQuoted", backtickQuotedQuery, []lexerOption{WithDBMS(DBMSMySQL)}},
 	}
 	obfuscator := NewObfuscator(
 		WithReplaceDigits(true),
@@ -141,14 +142,15 @@ ORDER BY
 		WithKeepSQLAlias(false),
 		WithUppercaseKeywords(true),
 		WithRemoveSpaceBetweenParentheses(true),
+		WithKeepIdentifierQuotation(true),
 	)
 
 	for _, bm := range benchmarks {
 		b.Run(bm.name+"/"+strconv.Itoa(len(bm.query)), func(b *testing.B) {
 			b.ResetTimer()
 			b.ReportAllocs()
 			for i := 0; i < b.N; i++ {
-				_, _, err := ObfuscateAndNormalize(bm.query, obfuscator, normalizer)
+				_, _, err := ObfuscateAndNormalize(bm.query, obfuscator, normalizer, bm.lexerOptions...)
 				if err != nil {
 					b.Fatal(err)
 				}