Fix(mysql): parse column prefix in index / pk defn. correctly (#2197)

* Fix(mysql): parse column prefix in index / pk defn. correctly

* Cleanup

sqlglot/dialects/mysql.py

@@ -358,6 +358,15 @@ class Parser(parser.Parser):
 
         LOG_DEFAULTS_TO_LN = True
 
+        def _parse_primary_key_part(self) -> t.Optional[exp.Expression]:
+            this = self._parse_id_var()
+            if not self._match(TokenType.L_PAREN):
+                return this
+
+            expression = self._parse_number()
+            self._match_r_paren()
+            return self.expression(exp.ColumnPrefix, this=this, expression=expression)
+
         def _parse_index_constraint(
             self, kind: t.Optional[str] = None
         ) -> exp.IndexColumnConstraint:

sqlglot/expressions.py

@@ -1538,6 +1538,10 @@ class ForeignKey(Expression):
     }
 
 
+class ColumnPrefix(Expression):
+    arg_types = {"this": True, "expression": True}
+
+
 class PrimaryKey(Expression):
     arg_types = {"expressions": True, "options": False}
 

sqlglot/generator.py

@@ -2771,6 +2771,9 @@ def comprehension_sql(self, expression: exp.Comprehension) -> str:
         condition = f" IF {condition}" if condition else ""
         return f"{this} FOR {expr} IN {iterator}{condition}"
 
+    def columnprefix_sql(self, expression: exp.ColumnPrefix) -> str:
+        return f"{self.sql(expression, 'this')}({self.sql(expression, 'expression')})"
+
 
 def cached_generator(
     cache: t.Optional[t.Dict[int, str]] = None

sqlglot/parser.py

@@ -3890,6 +3890,9 @@ def _parse_foreign_key(self) -> exp.ForeignKey:
             exp.ForeignKey, expressions=expressions, reference=reference, **options  # type: ignore
         )
 
+    def _parse_primary_key_part(self) -> t.Optional[exp.Expression]:
+        return self._parse_field()
+
     def _parse_primary_key(
         self, wrapped_optional: bool = False, in_props: bool = False
     ) -> exp.PrimaryKeyColumnConstraint | exp.PrimaryKey:
@@ -3901,7 +3904,9 @@ def _parse_primary_key(
         if not in_props and not self._match(TokenType.L_PAREN, advance=False):
             return self.expression(exp.PrimaryKeyColumnConstraint, desc=desc)
 
-        expressions = self._parse_wrapped_csv(self._parse_field, optional=wrapped_optional)
+        expressions = self._parse_wrapped_csv(
+            self._parse_primary_key_part, optional=wrapped_optional
+        )
         options = self._parse_key_constraint_options()
         return self.expression(exp.PrimaryKey, expressions=expressions, options=options)
 

tests/dialects/test_mysql.py

@@ -26,6 +26,9 @@ def test_ddl(self):
         self.validate_identity("CREATE TABLE foo (a BIGINT, INDEX USING BTREE (b))")
         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(
+            "CREATE TABLE `x` (`username` VARCHAR(200), PRIMARY KEY (`username`(16)))"
+        )
         self.validate_identity(
             "UPDATE items SET items.price = 0 WHERE items.id >= 5 ORDER BY items.id LIMIT 10"
         )