Feat(mysql): support ADD INDEX/KEY/UNIQUE in ALTER TABLE (#3621)

tobymao · Jun 10, 2024 · ee9b01d · ee9b01d
1 parent 05fe847
commit ee9b01d
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 5 deletions.
diff --git a/sqlglot/dialects/mysql.py b/sqlglot/dialects/mysql.py
@@ -202,6 +202,7 @@ class Tokenizer(tokens.Tokenizer):
             "CHARSET": TokenType.CHARACTER_SET,
             "FORCE": TokenType.FORCE,
             "IGNORE": TokenType.IGNORE,
+            "KEY": TokenType.KEY,
             "LOCK TABLES": TokenType.COMMAND,
             "LONGBLOB": TokenType.LONGBLOB,
             "LONGTEXT": TokenType.LONGTEXT,

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -1138,7 +1138,14 @@ class Parser(metaclass=_Parser):
 
     FETCH_TOKENS = ID_VAR_TOKENS - {TokenType.ROW, TokenType.ROWS, TokenType.PERCENT}
 
-    ADD_CONSTRAINT_TOKENS = {TokenType.CONSTRAINT, TokenType.PRIMARY_KEY, TokenType.FOREIGN_KEY}
+    ADD_CONSTRAINT_TOKENS = {
+        TokenType.CONSTRAINT,
+        TokenType.FOREIGN_KEY,
+        TokenType.INDEX,
+        TokenType.KEY,
+        TokenType.PRIMARY_KEY,
+        TokenType.UNIQUE,
+    }
 
     DISTINCT_TOKENS = {TokenType.DISTINCT}
 
@@ -3242,7 +3249,7 @@ def _parse_table_hints(self) -> t.Optional[t.List[exp.Expression]]:
             while self._match_set(self.TABLE_INDEX_HINT_TOKENS):
                 hint = exp.IndexTableHint(this=self._prev.text.upper())
 
-                self._match_texts(("INDEX", "KEY"))
+                self._match_set((TokenType.INDEX, TokenType.KEY))
                 if self._match(TokenType.FOR):
                     hint.set("target", self._advance_any() and self._prev.text.upper())
 

diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
@@ -287,6 +287,7 @@ class TokenType(AutoName):
     JOIN = auto()
     JOIN_MARKER = auto()
     KEEP = auto()
+    KEY = auto()
     KILL = auto()
     LANGUAGE = auto()
     LATERAL = auto()

diff --git a/tests/dialects/test_mysql.py b/tests/dialects/test_mysql.py
@@ -21,6 +21,9 @@ def test_ddl(self):
         self.validate_identity("CREATE TABLE foo (a BIGINT, FULLTEXT INDEX (b))")
         self.validate_identity("CREATE TABLE foo (a BIGINT, SPATIAL INDEX (b))")
         self.validate_identity("ALTER TABLE t1 ADD COLUMN x INT, ALGORITHM=INPLACE, LOCK=EXCLUSIVE")
+        self.validate_identity("ALTER TABLE t ADD INDEX `i` (`c`)")
+        self.validate_identity("ALTER TABLE t ADD UNIQUE `i` (`c`)")
+        self.validate_identity("ALTER TABLE test_table MODIFY COLUMN test_column LONGTEXT")
         self.validate_identity(
             "CREATE TABLE `oauth_consumer` (`key` VARCHAR(32) NOT NULL, UNIQUE `OAUTH_CONSUMER_KEY` (`key`))"
         )
@@ -60,6 +63,10 @@ def test_ddl(self):
         self.validate_identity(
             "CREATE OR REPLACE VIEW my_view AS SELECT column1 AS `boo`, column2 AS `foo` FROM my_table WHERE column3 = 'some_value' UNION SELECT q.* FROM fruits_table, JSON_TABLE(Fruits, '$[*]' COLUMNS(id VARCHAR(255) PATH '$.$id', value VARCHAR(255) PATH '$.value')) AS q",
         )
+        self.validate_identity(
+            "ALTER TABLE t ADD KEY `i` (`c`)",
+            "ALTER TABLE t ADD INDEX `i` (`c`)",
+        )
         self.validate_identity(
             "CREATE TABLE `foo` (`id` char(36) NOT NULL DEFAULT (uuid()), PRIMARY KEY (`id`), UNIQUE KEY `id` (`id`))",
             "CREATE TABLE `foo` (`id` CHAR(36) NOT NULL DEFAULT (UUID()), PRIMARY KEY (`id`), UNIQUE `id` (`id`))",
@@ -76,9 +83,6 @@ def test_ddl(self):
             "ALTER TABLE test_table ALTER COLUMN test_column SET DATA TYPE LONGTEXT",
             "ALTER TABLE test_table MODIFY COLUMN test_column LONGTEXT",
         )
-        self.validate_identity(
-            "ALTER TABLE test_table MODIFY COLUMN test_column LONGTEXT",
-        )
         self.validate_identity(
             "CREATE TABLE t (c DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP) DEFAULT CHARSET=utf8 ROW_FORMAT=DYNAMIC",
             "CREATE TABLE t (c DATETIME DEFAULT CURRENT_TIMESTAMP() ON UPDATE CURRENT_TIMESTAMP()) DEFAULT CHARACTER SET=utf8 ROW_FORMAT=DYNAMIC",