Fix(tsql): support adding multiple columns with ALTER TABLE (#2140)

* Support adding multiple columns with ALTER TABLE in tsql * Correctly parse tsql ALTER TABLE with multiple added columns * Fix styling
tobymao · Aug 31, 2023 · bd96d0c · bd96d0c
1 parent 2ad559d
commit bd96d0c
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 8 deletions.
diff --git a/sqlglot/dialects/tsql.py b/sqlglot/dialects/tsql.py
@@ -399,6 +399,8 @@ class Parser(parser.Parser):
 
         CONCAT_NULL_OUTPUTS_STRING = True
 
+        ALTER_TABLE_ADD_COLUMN_KEYWORD = False
+
         def _parse_projections(self) -> t.List[exp.Expression]:
             """
             T-SQL supports the syntax alias = expression in the SELECT's projection list,

diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -2271,11 +2271,14 @@ def altertable_sql(self, expression: exp.AlterTable) -> str:
         actions = expression.args["actions"]
 
         if isinstance(actions[0], exp.ColumnDef):
-            actions = self.expressions(
-                expression,
-                key="actions",
-                prefix="ADD COLUMN " if self.ALTER_TABLE_ADD_COLUMN_KEYWORD else "ADD ",
-            )
+            if self.ALTER_TABLE_ADD_COLUMN_KEYWORD:
+                actions = self.expressions(
+                    expression,
+                    key="actions",
+                    prefix="ADD COLUMN ",
+                )
+            else:
+                actions = f"ADD {self.expressions(expression, key='actions')}"
         elif isinstance(actions[0], exp.Schema):
             actions = self.expressions(expression, key="actions", prefix="ADD COLUMNS ")
         elif isinstance(actions[0], exp.Delete):

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -852,6 +852,9 @@ class Parser(metaclass=_Parser):
 
     SUPPORTS_USER_DEFINED_TYPES = True
 
+    # Whether or not ADD is present for each column added by ALTER TABLE
+    ALTER_TABLE_ADD_COLUMN_KEYWORD = True
+
     __slots__ = (
         "error_level",
         "error_message_context",
@@ -4708,6 +4711,9 @@ def _parse_alter_table_add(self) -> t.List[exp.Expression]:
             return self._parse_csv(self._parse_add_constraint)
 
         self._retreat(index)
+        if not self.ALTER_TABLE_ADD_COLUMN_KEYWORD and self._match_text_seq("ADD"):
+            return self._parse_csv(self._parse_field_def)
+
         return self._parse_csv(self._parse_add_column)
 
     def _parse_alter_table_alter(self) -> exp.AlterColumn:

diff --git a/tests/dialects/test_tsql.py b/tests/dialects/test_tsql.py
@@ -488,9 +488,14 @@ def test_ddl(self):
         )
 
         self.validate_all(
-            "ALTER TABLE a ADD b INTEGER",
-            read={"": "ALTER TABLE a ADD COLUMN b INT"},
-            write={"tsql": "ALTER TABLE a ADD b INTEGER"},
+            "ALTER TABLE a ADD b INTEGER, c INTEGER",
+            read={
+                "": "ALTER TABLE a ADD COLUMN b INT, ADD COLUMN c INT",
+            },
+            write={
+                "": "ALTER TABLE a ADD COLUMN b INT, ADD COLUMN c INT",
+                "tsql": "ALTER TABLE a ADD b INTEGER, c INTEGER",
+            },
         )
 
         self.validate_all(