Fix bug where table name is not properly collected on drop statement

DataDog · May 15, 2024 · 77f9b0d · 77f9b0d
1 parent a0ba802
commit 77f9b0d
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 0 deletions.
diff --git a/normalizer_test.go b/normalizer_test.go
@@ -357,6 +357,17 @@ multiline comment */
 				Size:       11,
 			},
 		},
+		{
+			input:    `DROP TABLE IF EXISTS users`,
+			expected: `DROP TABLE IF EXISTS users`,
+			statementMetadata: StatementMetadata{
+				Tables:     []string{"users"},
+				Comments:   []string{},
+				Commands:   []string{"DROP"},
+				Procedures: []string{},
+				Size:       9,
+			},
+		},
 	}
 
 	normalizer := NewNormalizer(

diff --git a/sqllexer_test.go b/sqllexer_test.go
@@ -632,6 +632,21 @@ func TestLexer(t *testing.T) {
 			},
 			lexerOpts: []lexerOption{WithDBMS(DBMSMySQL)},
 		},
+		{
+			name:  "drop table if exists",
+			input: `DROP TABLE IF EXISTS users`,
+			expected: []Token{
+				{IDENT, "DROP"},
+				{WS, " "},
+				{IDENT, "TABLE"},
+				{WS, " "},
+				{IDENT, "IF"},
+				{WS, " "},
+				{IDENT, "EXISTS"},
+				{WS, " "},
+				{IDENT, "users"},
+			},
+		},
 	}
 
 	for _, tt := range tests {

diff --git a/sqllexer_utils.go b/sqllexer_utils.go
@@ -49,6 +49,7 @@ var tableIndicators = map[string]bool{
 	"INTO":          true,
 	"UPDATE":        true,
 	"TABLE":         true,
+	"EXISTS":        true, // Drop Table If Exists
 	"STRAIGHT_JOIN": true, // MySQL
 	"CLONE":         true, // Snowflake
 }
@@ -145,6 +146,7 @@ var keywords = map[string]bool{
 	"OFFSET":     true,
 	"OF":         true,
 	"SKIP":       true,
+	"IF":         true,
 }
 
 func isWhitespace(ch rune) bool {