remove quotes from tables in metadata

DataDog · Dec 19, 2024 · 01a7786 · 01a7786
1 parent 31d6a30
commit 01a7786
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 20 deletions.
diff --git a/normalizer.go b/normalizer.go
@@ -165,27 +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 {
 			if !n.config.KeepIdentifierQuotation {
-				tokenVal = trimQuotes(tokenVal, tokenVal[0:1], tokenVal[len(tokenVal)-1:])
-				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)
 		}
 	}
 }

diff --git a/normalizer_test.go b/normalizer_test.go
@@ -756,11 +756,11 @@ func TestNormalizeDeobfuscatedSQL(t *testing.T) {
 			input:    `SELECT * FROM "users" WHERE id = ?`,
 			expected: `SELECT * FROM "users" WHERE id = ?`,
 			statementMetadata: StatementMetadata{
-				Tables:     []string{`"users"`},
+				Tables:     []string{`users`},
 				Comments:   []string{},
 				Commands:   []string{"SELECT"},
 				Procedures: []string{},
-				Size:       13,
+				Size:       11,
 			},
 			normalizationConfig: &normalizerConfig{
 				CollectComments:         true,
@@ -791,11 +791,11 @@ 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{},
-				Size:       22,
+				Size:       18,
 			},
 			normalizationConfig: &normalizerConfig{
 				CollectComments:         true,
@@ -829,11 +829,11 @@ 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{},
-				Size:       22,
+				Size:       18,
 			},
 			normalizationConfig: &normalizerConfig{
 				CollectComments:         true,
@@ -870,11 +870,11 @@ 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{},
-				Size:       22,
+				Size:       18,
 			},
 			normalizationConfig: &normalizerConfig{
 				CollectComments:         true,
@@ -891,11 +891,11 @@ func TestNormalizeDeobfuscatedSQL(t *testing.T) {
 			input:    `SELECT * FROM [public].[my users] WHERE id = ?`,
 			expected: `SELECT * FROM [public].[my users] WHERE id = ?`,
 			statementMetadata: StatementMetadata{
-				Tables:     []string{`[public].[my users]`},
+				Tables:     []string{`public.my users`},
 				Comments:   []string{},
 				Commands:   []string{"SELECT"},
 				Procedures: []string{},
-				Size:       25,
+				Size:       21,
 			},
 			normalizationConfig: &normalizerConfig{
 				CollectComments:         true,