fix(mysql): Parse REPLACE statement as Command (#3425)

tobymao · May 8, 2024 · e8014e2 · e8014e2
1 parent 5d4f148
commit e8014e2
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 7 deletions.
diff --git a/sqlglot/dialects/mysql.py b/sqlglot/dialects/mysql.py
@@ -251,7 +251,7 @@ class Tokenizer(tokens.Tokenizer):
             "@@": TokenType.SESSION_PARAMETER,
         }
 
-        COMMANDS = tokens.Tokenizer.COMMANDS - {TokenType.SHOW}
+        COMMANDS = {*tokens.Tokenizer.COMMANDS, TokenType.REPLACE} - {TokenType.SHOW}
 
     class Parser(parser.Parser):
         FUNC_TOKENS = {

diff --git a/sqlglot/dialects/tsql.py b/sqlglot/dialects/tsql.py
@@ -474,6 +474,8 @@ class Tokenizer(tokens.Tokenizer):
             "OPTION": TokenType.OPTION,
         }
 
+        COMMANDS = {*tokens.Tokenizer.COMMANDS, TokenType.END}
+
     class Parser(parser.Parser):
         SET_REQUIRES_ASSIGNMENT_DELIMITER = False
         LOG_DEFAULTS_TO_LN = True
@@ -524,11 +526,6 @@ class Parser(parser.Parser):
             *parser.Parser.TYPE_TOKENS,
         }
 
-        STATEMENT_PARSERS = {
-            **parser.Parser.STATEMENT_PARSERS,
-            TokenType.END: lambda self: self._parse_command(),
-        }
-
         def _parse_options(self) -> t.Optional[t.List[exp.Expression]]:
             if not self._match(TokenType.OPTION):
                 return None

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -1471,7 +1471,7 @@ def _parse_statement(self) -> t.Optional[exp.Expression]:
         if self._match_set(self.STATEMENT_PARSERS):
             return self.STATEMENT_PARSERS[self._prev.token_type](self)
 
-        if self._match_set(Tokenizer.COMMANDS):
+        if self._match_set(self.dialect.tokenizer.COMMANDS):
             return self._parse_command()
 
         expression = self._parse_expression()

diff --git a/tests/dialects/test_mysql.py b/tests/dialects/test_mysql.py
@@ -148,6 +148,9 @@ def test_identity(self):
         self.validate_identity(
             "SELECT * FROM t1, t2, t3 FOR SHARE OF t1 NOWAIT FOR UPDATE OF t2, t3 SKIP LOCKED"
         )
+        self.validate_identity(
+            "REPLACE INTO table SELECT id FROM table2 WHERE cnt > 100", check_command_warning=True
+        )
         self.validate_identity(
             """SELECT * FROM foo WHERE 3 MEMBER OF(info->'$.value')""",
             """SELECT * FROM foo WHERE 3 MEMBER OF(JSON_EXTRACT(info, '$.value'))""",