feat(clickhouse): add support for partition expression (#3428)

* feat(clickhouse): add support for partition expression

* feat(clickhouse): add support for partition expression

---------

Co-authored-by: Rustem Galiev <[email protected]>

Author: GaliFFun

sqlglot/dialects/clickhouse.py

@@ -578,6 +578,21 @@ def _parse_index_constraint(
                 granularity=granularity,
             )
 
+        def _parse_partition(self) -> t.Optional[exp.Partition]:
+            # https://clickhouse.com/docs/en/sql-reference/statements/alter/partition#how-to-set-partition-expression
+            if not self._match(TokenType.PARTITION):
+                return None
+
+            if self._match_text_seq("ID"):
+                # Corresponds to the PARTITION ID <string_value> syntax
+                expressions: t.List[exp.Expression] = [
+                    self.expression(exp.PartitionId, this=self._parse_string())
+                ]
+            else:
+                expressions = self._parse_expressions()
+
+            return self.expression(exp.Partition, expressions=expressions)
+
     class Generator(generator.Generator):
         QUERY_HINTS = False
         STRUCT_DELIMITER = ("(", ")")
@@ -828,3 +843,9 @@ def indexcolumnconstraint_sql(self, expression: exp.IndexColumnConstraint) -> st
             granularity = f" GRANULARITY {granularity}" if granularity else ""
 
             return f"INDEX{this}{expr}{index_type}{granularity}"
+
+        def partition_sql(self, expression: exp.Partition) -> str:
+            return f"PARTITION {self.expressions(expression, flat=True)}"
+
+        def partitionid_sql(self, expression: exp.PartitionId) -> str:
+            return f"ID {self.sql(expression.this)}"

sqlglot/expressions.py

@@ -2177,6 +2177,11 @@ class PartitionRange(Expression):
     arg_types = {"this": True, "expression": True}
 
 
+# https://clickhouse.com/docs/en/sql-reference/statements/alter/partition#how-to-set-partition-expression
+class PartitionId(Expression):
+    pass
+
+
 class Fetch(Expression):
     arg_types = {
         "direction": False,

tests/dialects/test_clickhouse.py

@@ -425,6 +425,13 @@ def test_clickhouse(self):
             },
         )
 
+        self.validate_identity("ALTER TABLE visits DROP PARTITION 201901")
+        self.validate_identity("ALTER TABLE visits DROP PARTITION ALL")
+        self.validate_identity(
+            "ALTER TABLE visits DROP PARTITION tuple(toYYYYMM(toDate('2019-01-25')))"
+        )
+        self.validate_identity("ALTER TABLE visits DROP PARTITION ID '201901'")
+
     def test_cte(self):
         self.validate_identity("WITH 'x' AS foo SELECT foo")
         self.validate_identity("WITH ['c'] AS field_names SELECT field_names")