Repository: tobymao/sqlglot
Pull request: Allow functions after IN
fixes #554 

Problem statement:
IN operation doesn't parse with the clickhouse dialect
clickhouse:

```
albatross :) select 'a' in mapKeys(map('a', 1, 'b', 2));

SELECT 'a' IN mapKeys(map('a', 1, 'b', 2))

Query id: 29c42995-45a2-491c-8b36-cf0e51b2b39b

┌─in('a', mapKeys(map('a', 1, 'b', 2)))─┐
│                                     1 │
└───────────────────────────────────────┘

1 row in set. Elapsed: 0.001 sec.
```

sqlglot:

```
In [1]: import sqlglot as sg

In [2]: sg.parse_one("'a' IN mapKeys(map('a', 1, 'b', 2))", read="clickhouse")

...

File /nix/store/gi2hjdx10z60l2q3m5zld3y2v273fg2k-python3-3.10.6-env/lib/python3.10/site-packages/sqlglot/parser.py:563, in Parser.check_errors(self)
    561         logger.error(str(error))
    562 elif self.error_level == ErrorLevel.RAISE and self.errors:
--> 563     raise ParseError(concat_errors(self.errors, self.max_errors))

ParseError: Expecting (. Line 1, Col: 8.
  'a' IN mapKeys(map('a', 1, 'b', 2))

Expecting ). Line 1, Col: 35.
  'a' IN mapKeys(map('a', 1, 'b', 2))
```

Hints from issue discussion:

Commit message: Allow functions after IN

Patch to review:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 39ce0af0ca..99c06140b7 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -2074,7 +2074,7 @@ class Distinct(Expression):
 
 
 class In(Predicate):
-    arg_types = {"this": True, "expressions": False, "query": False, "unnest": False}
+    arg_types = {"this": True, "expressions": False, "query": False, "unnest": False, "function": False}
 
 
 class TimeUnit(Expression):
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 0bd7175545..a8acf42639 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -945,10 +945,13 @@ def if_sql(self, expression):
     def in_sql(self, expression):
         query = expression.args.get("query")
         unnest = expression.args.get("unnest")
+        function = expression.args.get("function")
         if query:
             in_sql = self.wrap(query)
         elif unnest:
             in_sql = self.in_unnest_op(unnest)
+        elif function:
+            in_sql = self.sql(function)
         else:
             in_sql = f"({self.expressions(expression, flat=True)})"
         return f"{self.sql(expression, 'this')} IN {in_sql}"
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 1323ae0425..6f0dc94ad7 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -1504,8 +1504,7 @@ def _parse_in(self, this):
         unnest = self._parse_unnest()
         if unnest:
             this = self.expression(exp.In, this=this, unnest=unnest)
-        else:
-            self._match_l_paren()
+        elif self._match(TokenType.L_PAREN):
             expressions = self._parse_csv(self._parse_select_or_expression)
 
             if len(expressions) == 1 and isinstance(expressions[0], exp.Subqueryable):
@@ -1514,6 +1513,9 @@ def _parse_in(self, this):
                 this = self.expression(exp.In, this=this, expressions=expressions)
 
             self._match_r_paren()
+        else:
+            this = self.expression(exp.In, this=this, function=self._parse_field())
+
         return this
 
     def _parse_between(self, this):
diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
index 725b27c8bd..cf51acb991 100644
--- a/tests/dialects/test_clickhouse.py
+++ b/tests/dialects/test_clickhouse.py
@@ -8,6 +8,7 @@ def test_clickhouse(self):
         self.validate_identity("dictGet(x, 'y')")
         self.validate_identity("SELECT * FROM x FINAL")
         self.validate_identity("SELECT * FROM x AS y FINAL")
+        self.validate_identity("'a' IN mapKeys(map('a', 1, 'b', 2))")
 
         self.validate_all(
             "SELECT fname, lname, age FROM person ORDER BY age DESC NULLS FIRST, fname ASC NULLS LAST, lname",


Code review comments:
diff_hunk: @@ -2074,7 +2074,7 @@ class Distinct(Expression):
 
 
 class In(Predicate):
-    arg_types = {"this": True, "expressions": False, "query": False, "unnest": False}
+    arg_types = {"this": True, "expressions": False, "query": False, "unnest": False, "function": False}
path: sqlglot/expressions.py
text: @user1:
rename function to expression

@user1:
or field