fix: alter table add columns closes #3835

tobymao · tobymao · commit ae95c18f636d · 2024-07-25T11:24:09.000-07:00
diff --git a/sqlglot/dialects/spark2.py b/sqlglot/dialects/spark2.py
@@ -165,9 +165,6 @@ class Parser(Hive.Parser):
             "SHUFFLE_REPLICATE_NL": lambda self: self._parse_join_hint("SHUFFLE_REPLICATE_NL"),
         }
 
-        def _parse_add_column(self) -> t.Optional[exp.Expression]:
-            return self._match_text_seq("ADD", "COLUMNS") and self._parse_schema()
-
         def _parse_drop_column(self) -> t.Optional[exp.Drop | exp.Command]:
             return self._match_text_seq("DROP", "COLUMNS") and self.expression(
                 exp.Drop, this=self._parse_schema(), kind="COLUMNS"
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -3139,6 +3139,7 @@ def add_column_sql(self, expression: exp.AlterTable) -> str:
                 expression,
                 key="actions",
                 prefix="ADD COLUMN ",
+                skip_first=True,
             )
         return f"ADD {self.expressions(expression, key='actions', flat=True)}"
 
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -6305,6 +6305,13 @@ def _parse_alter_table_add(self) -> t.List[exp.Expression]:
         self._retreat(index)
         if not self.ALTER_TABLE_ADD_REQUIRED_FOR_EACH_COLUMN and self._match_text_seq("ADD"):
             return self._parse_wrapped_csv(self._parse_field_def, optional=True)
+
+        if self._match_text_seq("ADD", "COLUMNS"):
+            schema = self._parse_schema()
+            if schema:
+                return [schema]
+            return []
+
         return self._parse_wrapped_csv(self._parse_add_column, optional=True)
 
     def _parse_alter_table_alter(self) -> t.Optional[exp.Expression]:
diff --git a/tests/dialects/test_databricks.py b/tests/dialects/test_databricks.py
@@ -7,6 +7,7 @@ class TestDatabricks(Validator):
     dialect = "databricks"
 
     def test_databricks(self):
+        self.validate_identity("ALTER TABLE labels ADD COLUMN label_score FLOAT")
         self.validate_identity("DESCRIBE HISTORY a.b")
         self.validate_identity("DESCRIBE history.tbl")
         self.validate_identity("CREATE TABLE t (a STRUCT<c: MAP<STRING, STRING>>)")

Original file line number	Diff line number	Diff line change
`@@ -165,9 +165,6 @@ class Parser(Hive.Parser):`
`165`	`165`	`"SHUFFLE_REPLICATE_NL": lambda self: self._parse_join_hint("SHUFFLE_REPLICATE_NL"),`
`166`	`166`	`}`
`167`	`167`
`168`		`- def _parse_add_column(self) -> t.Optional[exp.Expression]:`
`169`		`- return self._match_text_seq("ADD", "COLUMNS") and self._parse_schema()`
`170`		`-`
`171`	`168`	`def _parse_drop_column(self) -> t.Optional[exp.Drop \| exp.Command]:`
`172`	`169`	`return self._match_text_seq("DROP", "COLUMNS") and self.expression(`
`173`	`170`	`exp.Drop, this=self._parse_schema(), kind="COLUMNS"`
Original file line number	Diff line number	Diff line change
`@@ -3139,6 +3139,7 @@ def add_column_sql(self, expression: exp.AlterTable) -> str:`
`3139`	`3139`	`expression,`
`3140`	`3140`	`key="actions",`
`3141`	`3141`	`prefix="ADD COLUMN ",`
	`3142`	`+ skip_first=True,`
`3142`	`3143`	`)`
`3143`	`3144`	`return f"ADD {self.expressions(expression, key='actions', flat=True)}"`
`3144`	`3145`