Fix(mysql): add support for index type in the UNIQUE KEY constraint (#…

…2211)
tobymao · Sep 13, 2023 · 14c1b3b · 14c1b3b
1 parent 416b341
commit 14c1b3b
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 9 deletions.
diff --git a/sqlglot/dialects/mysql.py b/sqlglot/dialects/mysql.py
@@ -373,7 +373,7 @@ def _parse_index_constraint(
                 self._match_texts({"INDEX", "KEY"})
 
             this = self._parse_id_var(any_token=False)
-            type_ = self._match(TokenType.USING) and self._advance_any() and self._prev.text
+            index_type = self._match(TokenType.USING) and self._advance_any() and self._prev.text
             schema = self._parse_schema()
 
             options = []
@@ -413,7 +413,7 @@ def _parse_index_constraint(
                 this=this,
                 schema=schema,
                 kind=kind,
-                type=type_,
+                index_type=index_type,
                 options=options,
             )
 

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -1321,7 +1321,13 @@ class GeneratedAsIdentityColumnConstraint(ColumnConstraintKind):
 
 # https://dev.mysql.com/doc/refman/8.0/en/create-table.html
 class IndexColumnConstraint(ColumnConstraintKind):
-    arg_types = {"this": False, "schema": True, "kind": False, "type": False, "options": False}
+    arg_types = {
+        "this": False,
+        "schema": True,
+        "kind": False,
+        "index_type": False,
+        "options": False,
+    }
 
 
 class InlineLengthColumnConstraint(ColumnConstraintKind):
@@ -1354,7 +1360,7 @@ class TitleColumnConstraint(ColumnConstraintKind):
 
 
 class UniqueColumnConstraint(ColumnConstraintKind):
-    arg_types = {"this": False}
+    arg_types = {"this": False, "index_type": False}
 
 
 class UppercaseColumnConstraint(ColumnConstraintKind):

diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -705,7 +705,9 @@ def primarykeycolumnconstraint_sql(self, expression: exp.PrimaryKeyColumnConstra
     def uniquecolumnconstraint_sql(self, expression: exp.UniqueColumnConstraint) -> str:
         this = self.sql(expression, "this")
         this = f" {this}" if this else ""
-        return f"UNIQUE{this}"
+        index_type = expression.args.get("index_type")
+        index_type = f" USING {index_type}" if index_type else ""
+        return f"UNIQUE{this}{index_type}"
 
     def createable_sql(self, expression: exp.Create, locations: t.DefaultDict) -> str:
         return self.sql(expression, "this")
@@ -2740,13 +2742,13 @@ def indexcolumnconstraint_sql(self, expression: exp.IndexColumnConstraint) -> st
         kind = f"{kind} INDEX" if kind else "INDEX"
         this = self.sql(expression, "this")
         this = f" {this}" if this else ""
-        type_ = self.sql(expression, "type")
-        type_ = f" USING {type_}" if type_ else ""
+        index_type = self.sql(expression, "index_type")
+        index_type = f" USING {index_type}" if index_type else ""
         schema = self.sql(expression, "schema")
         schema = f" {schema}" if schema else ""
         options = self.expressions(expression, key="options", sep=" ")
         options = f" {options}" if options else ""
-        return f"{kind}{this}{type_}{schema}{options}"
+        return f"{kind}{this}{index_type}{schema}{options}"
 
     def nvl2_sql(self, expression: exp.Nvl2) -> str:
         if self.NVL2_SUPPORTED:

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -3815,7 +3815,9 @@ def _parse_unnamed_constraint(
     def _parse_unique(self) -> exp.UniqueColumnConstraint:
         self._match_text_seq("KEY")
         return self.expression(
-            exp.UniqueColumnConstraint, this=self._parse_schema(self._parse_id_var(any_token=False))
+            exp.UniqueColumnConstraint,
+            this=self._parse_schema(self._parse_id_var(any_token=False)),
+            index_type=self._match(TokenType.USING) and self._advance_any() and self._prev.text,
         )
 
     def _parse_key_constraint_options(self) -> t.List[str]:

diff --git a/tests/dialects/test_mysql.py b/tests/dialects/test_mysql.py
@@ -19,6 +19,7 @@ def test_ddl(self):
             },
         )
 
+        self.validate_identity("CREATE TABLE foo (a BIGINT, UNIQUE (b) USING BTREE)")
         self.validate_identity("CREATE TABLE foo (id BIGINT)")
         self.validate_identity("CREATE TABLE 00f (1d BIGINT)")
         self.validate_identity("UPDATE items SET items.price = 0 WHERE items.id >= 5 LIMIT 10")