From f4d01f9ff01c52f7eae105325582a60e435eeae8 Mon Sep 17 00:00:00 2001
From: vaggelisd <daniasevangelos@gmail.com>
Date: Fri, 20 Sep 2024 15:53:35 +0300
Subject: [PATCH] feat(clickhouse): Add support for APPLY query modifier

---
 sqlglot/dialects/clickhouse.py    | 10 ++++++++++
 sqlglot/expressions.py            |  1 +
 sqlglot/generator.py              |  5 ++++-
 sqlglot/parser.py                 |  4 ++++
 tests/dialects/test_clickhouse.py |  6 ++++++
 5 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
index ec1984e94..19ab772ab 100644
--- a/sqlglot/dialects/clickhouse.py
+++ b/sqlglot/dialects/clickhouse.py
@@ -756,6 +756,16 @@ def _parse_projection_def(self) -> t.Optional[exp.ProjectionDef]:
         def _parse_constraint(self) -> t.Optional[exp.Expression]:
             return super()._parse_constraint() or self._parse_projection_def()
 
+        def _parse_alias(
+            self, this: t.Optional[exp.Expression], explicit: bool = False
+        ) -> t.Optional[exp.Expression]:
+            # In clickhouse "SELECT <expr> APPLY(...)" is a query modifier,
+            # so "APPLY" shouldn't be parsed as <expr>'s alias. However, "SELECT <expr> apply" is a valid alias
+            if self._match_pair(TokenType.APPLY, TokenType.L_PAREN, advance=False):
+                return this
+
+            return super()._parse_alias(this=this, explicit=explicit)
+
     class Generator(generator.Generator):
         QUERY_HINTS = False
         STRUCT_DELIMITER = ("(", ")")
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index d44b2c3de..21f337497 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -3103,6 +3103,7 @@ def isin(
     "settings": False,
     "format": False,
     "options": False,
+    "apply": False,
 }
 
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 80c862880..c5865d652 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -2491,9 +2491,12 @@ def select_sql(self, expression: exp.Select) -> str:
         # are the only dialects that use LIMIT_IS_TOP and both place DISTINCT first.
         top_distinct = f"{distinct}{hint}{top}" if self.LIMIT_IS_TOP else f"{top}{hint}{distinct}"
         expressions = f"{self.sep()}{expressions}" if expressions else expressions
+        apply = self.expressions(expression, key="apply", sep=" ")
+        apply = f"{self.sep()}{apply}" if apply else ""
+
         sql = self.query_modifiers(
             expression,
-            f"SELECT{top_distinct}{kind}{expressions}",
+            f"SELECT{top_distinct}{kind}{expressions}{apply}",
             self.sql(expression, "into", comment=False),
             self.sql(expression, "from", comment=False),
         )
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 045a65c3a..727689d1e 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -555,6 +555,7 @@ class Parser(metaclass=_Parser):
     TRIM_TYPES = {"LEADING", "TRAILING", "BOTH"}
 
     FUNC_TOKENS = {
+        TokenType.APPLY,
         TokenType.COLLATE,
         TokenType.COMMAND,
         TokenType.CURRENT_DATE,
@@ -2960,6 +2961,9 @@ def _parse_select(
             )
             this.comments = comments
 
+            while self._match_pair(TokenType.APPLY, TokenType.L_PAREN, advance=False):
+                this.append("apply", self._parse_function())
+
             into = self._parse_into()
             if into:
                 this.set("into", into)
diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
index d203141df..f35243afe 100644
--- a/tests/dialects/test_clickhouse.py
+++ b/tests/dialects/test_clickhouse.py
@@ -519,6 +519,12 @@ def test_clickhouse(self):
         self.validate_identity("SELECT TRIM(LEADING '(' FROM '(   Hello, world!   )')")
         self.validate_identity("current_timestamp").assert_is(exp.Column)
 
+        self.validate_identity("SELECT * APPLY(sum) FROM columns_transformers")
+        self.validate_identity("SELECT COLUMNS('[jk]') APPLY(toString) FROM columns_transformers")
+        self.validate_identity(
+            "SELECT COLUMNS('[jk]') APPLY(toString) APPLY(length) APPLY(max) FROM columns_transformers"
+        )
+
     def test_clickhouse_values(self):
         values = exp.select("*").from_(
             exp.values([exp.tuple_(1, 2, 3)], alias="subq", columns=["a", "b", "c"])