Feat(clickhouse): parse ternary operator (#1603)

* Feat(clickhouse): parse ternary operator * Formatting * Refactor * Fixup * Fixup
tobymao · May 12, 2023 · f585eef · f585eef
1 parent 6a62e18
commit f585eef
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 11 deletions.
diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
@@ -78,12 +78,33 @@ class Parser(parser.Parser):
             and self._parse_in(this, is_global=True),
         }
 
-        JOIN_KINDS = {*parser.Parser.JOIN_KINDS, TokenType.ANY, TokenType.ASOF}  # type: ignore
+        # The PLACEHOLDER entry is popped because 1) it doesn't affect Clickhouse (it corresponds to
+        # the postgres-specific JSONBContains parser) and 2) it makes parsing the ternary op simpler.
+        COLUMN_OPERATORS = parser.Parser.COLUMN_OPERATORS.copy()
+        COLUMN_OPERATORS.pop(TokenType.PLACEHOLDER)
 
-        TABLE_ALIAS_TOKENS = {*parser.Parser.TABLE_ALIAS_TOKENS} - {TokenType.ANY}  # type: ignore
+        JOIN_KINDS = {*parser.Parser.JOIN_KINDS, TokenType.ANY, TokenType.ASOF}
+
+        TABLE_ALIAS_TOKENS = {*parser.Parser.TABLE_ALIAS_TOKENS} - {TokenType.ANY}
 
         LOG_DEFAULTS_TO_LN = True
 
+        def _parse_expression(self, explicit_alias: bool = False) -> t.Optional[exp.Expression]:
+            return self._parse_alias(self._parse_ternary(), explicit=explicit_alias)
+
+        def _parse_ternary(self) -> t.Optional[exp.Expression]:
+            this = self._parse_conjunction()
+
+            if self._match(TokenType.PLACEHOLDER):
+                return self.expression(
+                    exp.If,
+                    this=this,
+                    true=self._parse_conjunction(),
+                    false=self._match(TokenType.COLON) and self._parse_conjunction(),
+                )
+
+            return this
+
         def _parse_in(
             self, this: t.Optional[exp.Expression], is_global: bool = False
         ) -> exp.Expression:

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -552,12 +552,12 @@ class Parser(metaclass=_Parser):
     RANGE_PARSERS = {
         TokenType.BETWEEN: lambda self, this: self._parse_between(this),
         TokenType.GLOB: binary_range_parser(exp.Glob),
-        TokenType.OVERLAPS: binary_range_parser(exp.Overlaps),
+        TokenType.ILIKE: binary_range_parser(exp.ILike),
         TokenType.IN: lambda self, this: self._parse_in(this),
+        TokenType.IRLIKE: binary_range_parser(exp.RegexpILike),
         TokenType.IS: lambda self, this: self._parse_is(this),
         TokenType.LIKE: binary_range_parser(exp.Like),
-        TokenType.ILIKE: binary_range_parser(exp.ILike),
-        TokenType.IRLIKE: binary_range_parser(exp.RegexpILike),
+        TokenType.OVERLAPS: binary_range_parser(exp.Overlaps),
         TokenType.RLIKE: binary_range_parser(exp.RegexpLike),
         TokenType.SIMILAR_TO: binary_range_parser(exp.SimilarTo),
     }
@@ -2677,8 +2677,8 @@ def _parse_set_operations(self, this: t.Optional[exp.Expression]) -> t.Optional[
             expression=self._parse_set_operations(self._parse_select(nested=True)),
         )
 
-    def _parse_expression(self) -> t.Optional[exp.Expression]:
-        return self._parse_alias(self._parse_conjunction())
+    def _parse_expression(self, explicit_alias: bool = False) -> t.Optional[exp.Expression]:
+        return self._parse_alias(self._parse_conjunction(), explicit=explicit_alias)
 
     def _parse_conjunction(self) -> t.Optional[exp.Expression]:
         return self._parse_tokens(self._parse_equality, self.CONJUNCTION)
@@ -2985,7 +2985,7 @@ def _parse_column(self) -> t.Optional[exp.Expression]:
                 field = self._parse_types()
                 if not field:
                     self.raise_error("Expected type")
-            elif op:
+            elif op and self._curr:
                 self._advance()
                 value = self._prev.text
                 field = (
@@ -3045,9 +3045,7 @@ def _parse_primary(self) -> t.Optional[exp.Expression]:
             if query:
                 expressions = [query]
             else:
-                expressions = self._parse_csv(
-                    lambda: self._parse_alias(self._parse_conjunction(), explicit=True)
-                )
+                expressions = self._parse_csv(lambda: self._parse_expression(explicit_alias=True))
 
             this = seq_get(expressions, 0)
             self._parse_query_modifiers(this)

diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
@@ -1,3 +1,4 @@
+from sqlglot import exp, parse_one
 from tests.dialects.test_dialect import Validator
 
 
@@ -68,6 +69,36 @@ def test_cte(self):
         self.validate_identity("WITH (SELECT foo) AS bar SELECT bar + 5")
         self.validate_identity("WITH test1 AS (SELECT i + 1, j + 1 FROM test1) SELECT * FROM test1")
 
+    def test_ternary(self):
+        self.validate_all("x ? 1 : 2", write={"clickhouse": "CASE WHEN x THEN 1 ELSE 2 END"})
+        self.validate_all(
+            "x AND FOO() > 3 + 2 ? 1 : 2",
+            write={"clickhouse": "CASE WHEN x AND FOO() > 3 + 2 THEN 1 ELSE 2 END"},
+        )
+        self.validate_all(
+            "x ? (y ? 1 : 2) : 3",
+            write={"clickhouse": "CASE WHEN x THEN (CASE WHEN y THEN 1 ELSE 2 END) ELSE 3 END"},
+        )
+        self.validate_all(
+            "x AND (foo() ? FALSE : TRUE) ? (y ? 1 : 2) : 3",
+            write={
+                "clickhouse": "CASE WHEN x AND (CASE WHEN foo() THEN FALSE ELSE TRUE END) THEN (CASE WHEN y THEN 1 ELSE 2 END) ELSE 3 END"
+            },
+        )
+
+        ternary = parse_one("x ? (y ? 1 : 2) : 3", read="clickhouse")
+
+        self.assertIsInstance(ternary, exp.If)
+        self.assertIsInstance(ternary.this, exp.Column)
+        self.assertIsInstance(ternary.args["true"], exp.Paren)
+        self.assertIsInstance(ternary.args["false"], exp.Literal)
+
+        nested_ternary = ternary.args["true"].this
+
+        self.assertIsInstance(nested_ternary.this, exp.Column)
+        self.assertIsInstance(nested_ternary.args["true"], exp.Literal)
+        self.assertIsInstance(nested_ternary.args["false"], exp.Literal)
+
     def test_signed_and_unsigned_types(self):
         data_types = [
             "UInt8",