Repository: tobymao/sqlglot
Pull request: Feat(spark): support ALTER ADD PARTITION
Fixes #5204

References:
- https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-partition#partition
- https://spark.apache.org/docs/latest/sql-ref-syntax-ddl-alter-table.html


Problem statement:
Alter table does not handle `ADD PARTITION`
I would expect this syntax to do a roundtrip successfully:
```
>>> sqlglot.parse_one("ALTER TABLE foo ADD PARTITION(event='click')").sql()
"ALTER TABLE foo PARTITION(event = 'click')"
```
but it drops the `ADD` portion. This is valid syntax in spark (where I ran into it) so it at least needs to be added to that dialect.

Hints from issue discussion:
What is the input dialect?
@georgesittas the code is a repro as is in a clean session. 

but if you want a dialect then as I said in the comment, use `spark`
Ah apologies, I missed it. Will take a look, thanks.

Commit message: Feat(spark): support ALTER ADD PARTITION

Patch to review:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 1003ae5faa..f1a927a7c1 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -4926,6 +4926,10 @@ class AddConstraint(Expression):
     arg_types = {"expressions": True}
 
 
+class AddPartition(Expression):
+    pass
+
+
 class AttachOption(Expression):
     arg_types = {"this": True, "expression": False}
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index a9e0545282..7c5f046395 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -3498,6 +3498,9 @@ def droppartition_sql(self, expression: exp.DropPartition) -> str:
     def addconstraint_sql(self, expression: exp.AddConstraint) -> str:
         return f"ADD {self.expressions(expression)}"
 
+    def addpartition_sql(self, expression: exp.AddPartition) -> str:
+        return f"ADD {self.sql(expression.this)}"
+
     def distinct_sql(self, expression: exp.Distinct) -> str:
         this = self.expressions(expression, flat=True)
 
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 04c1703aff..e67e954e58 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -7368,13 +7368,20 @@ def _parse_drop_partition(self, exists: t.Optional[bool] = None) -> exp.DropPart
         )
 
     def _parse_alter_table_add(self) -> t.List[exp.Expression]:
-        def _parse_add_column_or_constraint():
+        def _parse_add_alteration() -> t.Optional[exp.Expression]:
             self._match_text_seq("ADD")
             if self._match_set(self.ADD_CONSTRAINT_TOKENS, advance=False):
                 return self.expression(
                     exp.AddConstraint, expressions=self._parse_csv(self._parse_constraint)
                 )
-            return self._parse_add_column()
+
+            is_partition = self._match(TokenType.PARTITION, advance=False)
+            field = self._parse_add_column()
+
+            if is_partition:
+                return self.expression(exp.AddPartition, this=field)
+
+            return field
 
         if not self.dialect.ALTER_TABLE_ADD_REQUIRED_FOR_EACH_COLUMN or self._match_text_seq(
             "COLUMNS"
@@ -7383,7 +7390,7 @@ def _parse_add_column_or_constraint():
 
             return ensure_list(schema) if schema else self._parse_csv(self._parse_field_def)
 
-        return self._parse_csv(_parse_add_column_or_constraint)
+        return self._parse_csv(_parse_add_alteration)
 
     def _parse_alter_table_alter(self) -> t.Optional[exp.Expression]:
         if self._match_texts(self.ALTER_ALTER_PARSERS):
diff --git a/tests/dialects/test_spark.py b/tests/dialects/test_spark.py
index 2ff616cfe9..fec7d41311 100644
--- a/tests/dialects/test_spark.py
+++ b/tests/dialects/test_spark.py
@@ -247,6 +247,7 @@ def test_spark(self):
             "REFRESH TABLE t",
         )
 
+        self.validate_identity("ALTER TABLE foo ADD PARTITION(event = 'click')")
         self.validate_identity("IF(cond, foo AS bar, bla AS baz)")
         self.validate_identity("any_value(col, true)", "ANY_VALUE(col) IGNORE NULLS")
         self.validate_identity("first(col, true)", "FIRST(col) IGNORE NULLS")


Code review comments:
diff_hunk: @@ -7368,13 +7368,20 @@ def _parse_drop_partition(self, exists: t.Optional[bool] = None) -> exp.DropPart
         )
 
     def _parse_alter_table_add(self) -> t.List[exp.Expression]:
-        def _parse_add_column_or_constraint():
+        def _parse_add_alteration() -> t.Optional[exp.Expression]:
             self._match_text_seq("ADD")
             if self._match_set(self.ADD_CONSTRAINT_TOKENS, advance=False):
                 return self.expression(
                     exp.AddConstraint, expressions=self._parse_csv(self._parse_constraint)
                 )
-            return self._parse_add_column()
+
+            is_partition = self._match(TokenType.PARTITION, advance=False)
+            field = self._parse_add_column()
path: sqlglot/parser.py
text: @user1:
Looks like the docs also allow `[IF NOT EXISTS]` between `ADD` and `PARTITION` afaict, but feel free to ignore

@author:
Addressed [here](https://github.com/tobymao/sqlglot/pull/5208/commits/434b86cc6e932f27c499e5ef2665c831a29000a4), please take another look.