mypy

PostHog · Sep 25, 2024 · c7d47ab · c7d47ab
1 parent 0a3e220
commit c7d47ab
Show file tree

Hide file tree

Showing 22 changed files with 117 additions and 77 deletions.
diff --git a/hogql_parser/HogQLParser.cpp b/hogql_parser/HogQLParser.cpp
@@ -518,7 +518,7 @@ void hogqlparserParserInitialize() {
   	0,0,1069,1071,5,117,0,0,1070,1069,1,0,0,0,1070,1071,1,0,0,0,1071,1075,
   	1,0,0,0,1072,1073,5,131,0,0,1073,1075,5,150,0,0,1074,1047,1,0,0,0,1074,
   	1061,1,0,0,0,1074,1072,1,0,0,0,1075,1076,1,0,0,0,1076,1079,5,112,0,0,
-  	1077,1080,3,116,58,0,1078,1080,3,36,18,0,1079,1077,1,0,0,0,1079,1078,
+  	1077,1080,3,36,18,0,1078,1080,3,116,58,0,1079,1077,1,0,0,0,1079,1078,
   	1,0,0,0,1080,119,1,0,0,0,1081,1082,5,133,0,0,1082,1086,3,156,78,0,1083,
   	1085,3,122,61,0,1084,1083,1,0,0,0,1085,1088,1,0,0,0,1086,1084,1,0,0,0,
   	1086,1087,1,0,0,0,1087,1089,1,0,0,0,1088,1086,1,0,0,0,1089,1090,5,152,
@@ -578,7 +578,7 @@ void hogqlparserParserInitialize() {
   	1257,5,106,0,0,1254,1257,3,148,74,0,1255,1257,3,150,75,0,1256,1253,1,
   	0,0,0,1256,1254,1,0,0,0,1256,1255,1,0,0,0,1257,157,1,0,0,0,1258,1259,
   	3,162,81,0,1259,1260,5,123,0,0,1260,1261,3,144,72,0,1261,159,1,0,0,0,
-  	1262,1263,5,129,0,0,1263,1264,3,130,65,0,1264,1265,5,148,0,0,1265,161,
+  	1262,1263,5,129,0,0,1263,1264,3,116,58,0,1264,1265,5,148,0,0,1265,161,
   	1,0,0,0,1266,1269,5,111,0,0,1267,1269,3,164,82,0,1268,1266,1,0,0,0,1268,
   	1267,1,0,0,0,1269,163,1,0,0,0,1270,1274,5,143,0,0,1271,1273,3,166,83,
   	0,1272,1271,1,0,0,0,1273,1276,1,0,0,0,1274,1272,1,0,0,0,1274,1275,1,0,
@@ -8886,14 +8886,14 @@ tree::TerminalNode* HogQLParser::ColumnLambdaExprContext::RPAREN() {
   return getToken(HogQLParser::RPAREN, 0);
 }
 
-HogQLParser::ColumnExprContext* HogQLParser::ColumnLambdaExprContext::columnExpr() {
-  return getRuleContext<HogQLParser::ColumnExprContext>(0);
-}
-
 HogQLParser::BlockContext* HogQLParser::ColumnLambdaExprContext::block() {
   return getRuleContext<HogQLParser::BlockContext>(0);
 }
 
+HogQLParser::ColumnExprContext* HogQLParser::ColumnLambdaExprContext::columnExpr() {
+  return getRuleContext<HogQLParser::ColumnExprContext>(0);
+}
+
 std::vector<tree::TerminalNode *> HogQLParser::ColumnLambdaExprContext::COMMA() {
   return getTokens(HogQLParser::COMMA);
 }
@@ -9011,13 +9011,13 @@ HogQLParser::ColumnLambdaExprContext* HogQLParser::columnLambdaExpr() {
     switch (getInterpreter<atn::ParserATNSimulator>()->adaptivePredict(_input, 136, _ctx)) {
     case 1: {
       setState(1077);
-      columnExpr(0);
+      block();
       break;
     }
 
     case 2: {
       setState(1078);
-      block();
+      columnExpr(0);
       break;
     }
 
@@ -11599,8 +11599,8 @@ tree::TerminalNode* HogQLParser::PlaceholderContext::LBRACE() {
   return getToken(HogQLParser::LBRACE, 0);
 }
 
-HogQLParser::NestedIdentifierContext* HogQLParser::PlaceholderContext::nestedIdentifier() {
-  return getRuleContext<HogQLParser::NestedIdentifierContext>(0);
+HogQLParser::ColumnExprContext* HogQLParser::PlaceholderContext::columnExpr() {
+  return getRuleContext<HogQLParser::ColumnExprContext>(0);
 }
 
 tree::TerminalNode* HogQLParser::PlaceholderContext::RBRACE() {
@@ -11636,7 +11636,7 @@ HogQLParser::PlaceholderContext* HogQLParser::placeholder() {
     setState(1262);
     match(HogQLParser::LBRACE);
     setState(1263);
-    nestedIdentifier();
+    columnExpr(0);
     setState(1264);
     match(HogQLParser::RBRACE);
 

diff --git a/hogql_parser/HogQLParser.h b/hogql_parser/HogQLParser.h
@@ -1872,8 +1872,8 @@ class  HogQLParser : public antlr4::Parser {
     std::vector<IdentifierContext *> identifier();
     IdentifierContext* identifier(size_t i);
     antlr4::tree::TerminalNode *RPAREN();
-    ColumnExprContext *columnExpr();
     BlockContext *block();
+    ColumnExprContext *columnExpr();
     std::vector<antlr4::tree::TerminalNode *> COMMA();
     antlr4::tree::TerminalNode* COMMA(size_t i);
 
@@ -2407,7 +2407,7 @@ class  HogQLParser : public antlr4::Parser {
     PlaceholderContext(antlr4::ParserRuleContext *parent, size_t invokingState);
     virtual size_t getRuleIndex() const override;
     antlr4::tree::TerminalNode *LBRACE();
-    NestedIdentifierContext *nestedIdentifier();
+    ColumnExprContext *columnExpr();
     antlr4::tree::TerminalNode *RBRACE();
 
 

diff --git a/hogql_parser/HogQLParser.interp b/hogql_parser/HogQLParser.interp
diff --git a/hogql_parser/parser.cpp b/hogql_parser/parser.cpp
@@ -2570,11 +2570,7 @@ class HogQLParseTreeConverter : public HogQLParserBaseVisitor {
   }
 
   VISIT(Placeholder) {
-    auto nested_identifier_ctx = ctx->nestedIdentifier();
-    vector<string> nested =
-        nested_identifier_ctx ? any_cast<vector<string>>(visit(nested_identifier_ctx)) : vector<string>();
-
-    RETURN_NEW_AST_NODE("Placeholder", "{s:N}", "chain", X_PyList_FromStrings(nested));
+    RETURN_NEW_AST_NODE("Placeholder", "{s:N}", "expr", visitAsPyObject(ctx->columnExpr()));
   }
 
   VISIT_UNSUPPORTED(EnumValue)

diff --git a/hogql_parser/pyproject.toml b/hogql_parser/pyproject.toml
@@ -17,7 +17,7 @@ archs = [ # We could also build a universal wheel, but separate ones are lighter
     "arm64",
 ]
 before-build = [ # We need to install the libraries for each architecture separately
-    "brew uninstall --force boost antlr4-cpp-runtime",
+    "brew uninstall --ignore-dependencies --force boost antlr4-cpp-runtime",
     "brew fetch --force --bottle-tag=${ARCHFLAGS##'-arch '}_monterey boost antlr4-cpp-runtime",
     "brew install $(brew --cache --bottle-tag=${ARCHFLAGS##'-arch '}_monterey boost antlr4-cpp-runtime)",
 ]

diff --git a/hogql_parser/setup.py b/hogql_parser/setup.py
@@ -32,7 +32,7 @@
 
 setup(
     name="hogql_parser",
-    version="1.0.40",
+    version="1.0.44",
     url="https://github.com/PostHog/posthog/tree/master/hogql_parser",
     author="PostHog Inc.",
     author_email="[email protected]",

diff --git a/posthog/api/test/test_insight.py b/posthog/api/test/test_insight.py
@@ -2852,7 +2852,7 @@ def test_insight_trend_hogql_global_filters(self) -> None:
             )
             self.assertEqual(
                 response_placeholder.json(),
-                self.validation_error_response("Placeholders, such as {team_id}, are not supported in this context"),
+                self.validation_error_response("Unresolved placeholder: {team_id}"),
             )
 
     @also_test_with_materialized_columns(event_properties=["int_value"], person_properties=["fish"])

diff --git a/posthog/hogql/ast.py b/posthog/hogql/ast.py
@@ -698,11 +698,15 @@ class Field(Expr):
 
 @dataclass(kw_only=True)
 class Placeholder(Expr):
-    chain: list[str | int]
+    expr: Expr
+
+    @property
+    def chain(self) -> list[str | int] | None:
+        return self.expr.chain if isinstance(self.expr, Field) else None
 
     @property
-    def field(self):
-        return ".".join(str(chain) for chain in self.chain)
+    def field(self) -> str | None:
+        return ".".join(str(chain) for chain in self.chain) if self.chain else None
 
 
 @dataclass(kw_only=True)

diff --git a/posthog/hogql/grammar/HogQLParser.g4 b/posthog/hogql/grammar/HogQLParser.g4
@@ -215,7 +215,7 @@ columnLambdaExpr:
     |        identifier (COMMA identifier)* COMMA?
     | LPAREN RPAREN
     )
-    ARROW (columnExpr | block)
+    ARROW (block | columnExpr)
     ;
 
 
@@ -293,7 +293,7 @@ keywordForAlias
 alias: IDENTIFIER | keywordForAlias;  // |interval| can't be an alias, otherwise 'INTERVAL 1 SOMETHING' becomes ambiguous.
 identifier: IDENTIFIER | interval | keyword;
 enumValue: string EQ_SINGLE numberLiteral;
-placeholder: LBRACE nestedIdentifier RBRACE;
+placeholder: LBRACE columnExpr RBRACE;
 
 string: STRING_LITERAL | templateString;
 templateString : QUOTE_SINGLE_TEMPLATE stringContents* QUOTE_SINGLE ;

diff --git a/posthog/hogql/grammar/HogQLParser.interp b/posthog/hogql/grammar/HogQLParser.interp
diff --git a/posthog/hogql/grammar/HogQLParser.py b/posthog/hogql/grammar/HogQLParser.py
@@ -446,8 +446,8 @@ def serializedATN():
         1070,1,0,0,0,1068,1066,1,0,0,0,1069,1071,5,117,0,0,1070,1069,1,0,
         0,0,1070,1071,1,0,0,0,1071,1075,1,0,0,0,1072,1073,5,131,0,0,1073,
         1075,5,150,0,0,1074,1047,1,0,0,0,1074,1061,1,0,0,0,1074,1072,1,0,
-        0,0,1075,1076,1,0,0,0,1076,1079,5,112,0,0,1077,1080,3,116,58,0,1078,
-        1080,3,36,18,0,1079,1077,1,0,0,0,1079,1078,1,0,0,0,1080,119,1,0,
+        0,0,1075,1076,1,0,0,0,1076,1079,5,112,0,0,1077,1080,3,36,18,0,1078,
+        1080,3,116,58,0,1079,1077,1,0,0,0,1079,1078,1,0,0,0,1080,119,1,0,
         0,0,1081,1082,5,133,0,0,1082,1086,3,156,78,0,1083,1085,3,122,61,
         0,1084,1083,1,0,0,0,1085,1088,1,0,0,0,1086,1084,1,0,0,0,1086,1087,
         1,0,0,0,1087,1089,1,0,0,0,1088,1086,1,0,0,0,1089,1090,5,152,0,0,
@@ -511,7 +511,7 @@ def serializedATN():
         5,106,0,0,1254,1257,3,148,74,0,1255,1257,3,150,75,0,1256,1253,1,
         0,0,0,1256,1254,1,0,0,0,1256,1255,1,0,0,0,1257,157,1,0,0,0,1258,
         1259,3,162,81,0,1259,1260,5,123,0,0,1260,1261,3,144,72,0,1261,159,
-        1,0,0,0,1262,1263,5,129,0,0,1263,1264,3,130,65,0,1264,1265,5,148,
+        1,0,0,0,1262,1263,5,129,0,0,1263,1264,3,116,58,0,1264,1265,5,148,
         0,0,1265,161,1,0,0,0,1266,1269,5,111,0,0,1267,1269,3,164,82,0,1268,
         1266,1,0,0,0,1268,1267,1,0,0,0,1269,163,1,0,0,0,1270,1274,5,143,
         0,0,1271,1273,3,166,83,0,1272,1271,1,0,0,0,1273,1276,1,0,0,0,1274,
@@ -7663,14 +7663,14 @@ def identifier(self, i:int=None):
         def RPAREN(self):
             return self.getToken(HogQLParser.RPAREN, 0)
 
-        def columnExpr(self):
-            return self.getTypedRuleContext(HogQLParser.ColumnExprContext,0)
-
-
         def block(self):
             return self.getTypedRuleContext(HogQLParser.BlockContext,0)
 
 
+        def columnExpr(self):
+            return self.getTypedRuleContext(HogQLParser.ColumnExprContext,0)
+
+
         def COMMA(self, i:int=None):
             if i is None:
                 return self.getTokens(HogQLParser.COMMA)
@@ -7770,12 +7770,12 @@ def columnLambdaExpr(self):
             la_ = self._interp.adaptivePredict(self._input,136,self._ctx)
             if la_ == 1:
                 self.state = 1077
-                self.columnExpr(0)
+                self.block()
                 pass
 
             elif la_ == 2:
                 self.state = 1078
-                self.block()
+                self.columnExpr(0)
                 pass
 
 
@@ -9701,8 +9701,8 @@ def __init__(self, parser, parent:ParserRuleContext=None, invokingState:int=-1):
         def LBRACE(self):
             return self.getToken(HogQLParser.LBRACE, 0)
 
-        def nestedIdentifier(self):
-            return self.getTypedRuleContext(HogQLParser.NestedIdentifierContext,0)
+        def columnExpr(self):
+            return self.getTypedRuleContext(HogQLParser.ColumnExprContext,0)
 
 
         def RBRACE(self):
@@ -9729,7 +9729,7 @@ def placeholder(self):
             self.state = 1262
             self.match(HogQLParser.LBRACE)
             self.state = 1263
-            self.nestedIdentifier()
+            self.columnExpr(0)
             self.state = 1264
             self.match(HogQLParser.RBRACE)
         except RecognitionException as re:

diff --git a/posthog/hogql/parser.py b/posthog/hogql/parser.py
@@ -1122,8 +1122,7 @@ def visitHogqlxTagAttribute(self, ctx: HogQLParser.HogqlxTagAttributeContext):
             return ast.HogQLXAttribute(name=name, value=ast.Constant(value=True))
 
     def visitPlaceholder(self, ctx: HogQLParser.PlaceholderContext):
-        nested = self.visit(ctx.nestedIdentifier()) if ctx.nestedIdentifier() else []
-        return ast.Placeholder(chain=nested)
+        return ast.Placeholder(expr=self.visit(ctx.columnExpr()))
 
     def visitColumnExprTemplateString(self, ctx: HogQLParser.ColumnExprTemplateStringContext):
         return self.visit(ctx.templateString())

diff --git a/posthog/hogql/placeholders.py b/posthog/hogql/placeholders.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, cast
 
 from posthog.hogql import ast
 from posthog.hogql.errors import QueryError
@@ -9,7 +9,7 @@ def replace_placeholders(node: ast.Expr, placeholders: Optional[dict[str, ast.Ex
     return ReplacePlaceholders(placeholders).visit(node)
 
 
-def find_placeholders(node: ast.Expr) -> list[str]:
+def find_placeholders(node: ast.Expr) -> list[str | None]:
     finder = FindPlaceholders()
     finder.visit(node)
     return list(finder.found)
@@ -18,7 +18,7 @@ def find_placeholders(node: ast.Expr) -> list[str]:
 class FindPlaceholders(TraversingVisitor):
     def __init__(self):
         super().__init__()
-        self.found: set[str] = set()
+        self.found: set[str | None] = set()
 
     def visit_cte(self, node: ast.CTE):
         super().visit(node.expr)
@@ -32,9 +32,24 @@ def __init__(self, placeholders: Optional[dict[str, ast.Expr]]):
         super().__init__()
         self.placeholders = placeholders
 
+    # Lambdas in the form x -> {placeholder} need special handling
+    def visit_block(self, node):
+        if self.placeholders and len(node.declarations) == 1 and isinstance(node.declarations[0], ast.ExprStatement):
+            expr = cast(ast.ExprStatement, node.declarations[0]).expr
+            if isinstance(expr, ast.Field):
+                field = ".".join(str(f) for f in expr.chain)
+                if field in self.placeholders and self.placeholders[field] is not None:
+                    new_node = self.placeholders[field]
+                    new_node.start = node.start
+                    new_node.end = node.end
+                    return new_node
+        return super().visit_block(node)
+
     def visit_placeholder(self, node):
+        if node.field is None:
+            raise QueryError("Placeholder expressions are not yet supported")
         if not self.placeholders:
-            raise QueryError(f"Placeholders, such as {{{node.field}}}, are not supported in this context")
+            raise QueryError(f"Unresolved placeholder: {{{node.field}}}")
         if node.field in self.placeholders and self.placeholders[node.field] is not None:
             new_node = self.placeholders[node.field]
             new_node.start = node.start

diff --git a/posthog/hogql/printer.py b/posthog/hogql/printer.py
@@ -1165,7 +1165,7 @@ def visit_call(self, node: ast.Call):
             raise QueryError(f"Unsupported function call '{node.name}(...)'")
 
     def visit_placeholder(self, node: ast.Placeholder):
-        raise QueryError(f"Placeholders, such as {{{node.field}}}, are not supported in this context")
+        raise QueryError(f"Unresolved placeholder: {{{node.field}}}")
 
     def visit_alias(self, node: ast.Alias):
         # Skip hidden aliases completely.

diff --git a/posthog/hogql/query.py b/posthog/hogql/query.py
@@ -77,12 +77,14 @@ def execute_hogql_query(
                 f"Query contains 'filters' placeholder, yet filters are also provided as a standalone query parameter."
             )
         if "filters" in placeholders_in_query or any(
-            placeholder.startswith("filters.") for placeholder in placeholders_in_query
+            placeholder and placeholder.startswith("filters.") for placeholder in placeholders_in_query
         ):
             select_query = replace_filters(select_query, filters, team)
 
             leftover_placeholders: list[str] = []
             for placeholder in placeholders_in_query:
+                if placeholder is None:
+                    raise ValueError("Placeholder expressions are not yet supported")
                 if placeholder != "filters" and not placeholder.startswith("filters."):
                     leftover_placeholders.append(placeholder)
 

diff --git a/posthog/hogql/test/_test_parser.py b/posthog/hogql/test/_test_parser.py
@@ -735,7 +735,7 @@ def test_expr_with_ignored_sql_comment(self):
         def test_placeholders(self):
             self.assertEqual(
                 self._expr("{foo}"),
-                ast.Placeholder(chain=["foo"]),
+                ast.Placeholder(expr=ast.Field(chain=["foo"])),
             )
             self.assertEqual(
                 self._expr("{foo}", {"foo": ast.Constant(value="bar")}),
@@ -946,7 +946,7 @@ def test_select_from_placeholder(self):
                 self._select("select 1 from {placeholder}"),
                 ast.SelectQuery(
                     select=[ast.Constant(value=1)],
-                    select_from=ast.JoinExpr(table=ast.Placeholder(chain=["placeholder"])),
+                    select_from=ast.JoinExpr(table=ast.Placeholder(expr=ast.Field(chain=["placeholder"]))),
                 ),
             )
             self.assertEqual(
@@ -1336,7 +1336,7 @@ def test_select_placeholders(self):
                     where=ast.CompareOperation(
                         op=ast.CompareOperationOp.Eq,
                         left=ast.Constant(value=1),
-                        right=ast.Placeholder(chain=["hogql_val_1"]),
+                        right=ast.Placeholder(expr=ast.Field(chain=["hogql_val_1"])),
                     ),
                 ),
             )
@@ -1355,6 +1355,29 @@ def test_select_placeholders(self):
                 ),
             )
 
+        def test_placeholder_expressions(self):
+            actual = self._select("select 1 where 1 == {1 ? hogql_val_1 : hogql_val_2}")
+            expected = clear_locations(
+                ast.SelectQuery(
+                    select=[ast.Constant(value=1)],
+                    where=ast.CompareOperation(
+                        op=ast.CompareOperationOp.Eq,
+                        left=ast.Constant(value=1),
+                        right=ast.Placeholder(
+                            expr=ast.Call(
+                                name="if",
+                                args=[
+                                    ast.Constant(value=1),
+                                    ast.Field(chain=["hogql_val_1"]),
+                                    ast.Field(chain=["hogql_val_2"]),
+                                ],
+                            )
+                        ),
+                    ),
+                )
+            )
+            self.assertEqual(actual, expected)
+
         def test_select_union_all(self):
             self.assertEqual(
                 self._select("select 1 union all select 2 union all select 3"),