fix(hogql): Fixed being able to use placeholder nodes in CTEs (#20273)

* Fixed being able to use placeholder nodes in CTEs

* Added tests and fixed mypy

mypy-baseline.txt

@@ -21,11 +21,17 @@ posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incomp
 posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "JoinConstraint | None"; expected "AST"  [arg-type]
 posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "JoinExpr | None"; expected "AST"  [arg-type]
 posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "JoinExpr | None"; expected "AST"  [arg-type]
+posthog/hogql/visitor.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "CTE")  [assignment]
+posthog/hogql/visitor.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "CTE")  [assignment]
 posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "Expr | None"; expected "AST"  [arg-type]
 posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "Expr | None"; expected "AST"  [arg-type]
 posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "Expr | None"; expected "AST"  [arg-type]
+posthog/hogql/visitor.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "CTE")  [assignment]
+posthog/hogql/visitor.py:0: error: Incompatible types in assignment (expression has type "OrderExpr", variable has type "CTE")  [assignment]
+posthog/hogql/visitor.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "CTE")  [assignment]
 posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "Expr | None"; expected "AST"  [arg-type]
 posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "Expr | None"; expected "AST"  [arg-type]
+posthog/hogql/visitor.py:0: error: Incompatible types in assignment (expression has type "WindowExpr", variable has type "CTE")  [assignment]
 posthog/hogql/visitor.py:0: error: Incompatible types in assignment (expression has type "FieldAliasType", variable has type "BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType")  [assignment]
 posthog/hogql/visitor.py:0: error: Incompatible types in assignment (expression has type "Type", variable has type "BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType")  [assignment]
 posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "WindowFrameExpr | None"; expected "AST"  [arg-type]
@@ -301,8 +307,6 @@ posthog/hogql/transforms/in_cohort.py:0: error: Item "None" of "JoinConstraint |
 posthog/hogql/transforms/in_cohort.py:0: error: Item "Expr" of "Expr | Any" has no attribute "right"  [union-attr]
 posthog/hogql/transforms/in_cohort.py:0: error: List item 0 has incompatible type "SelectQueryType | None"; expected "SelectQueryType"  [list-item]
 posthog/hogql/transforms/in_cohort.py:0: error: List item 0 has incompatible type "SelectQueryType | None"; expected "SelectQueryType"  [list-item]
-posthog/api/organization.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
-posthog/api/organization.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
 posthog/hogql/printer.py:0: error: Argument 1 to "create_hogql_database" has incompatible type "int | None"; expected "int"  [arg-type]
 posthog/hogql/printer.py:0: error: List comprehension has incompatible type List[SelectQueryType | None]; expected List[SelectQueryType]  [misc]
 posthog/hogql/printer.py:0: error: Argument "stack" to "_Printer" has incompatible type "list[SelectQuery]"; expected "list[AST] | None"  [arg-type]
@@ -333,6 +337,8 @@ posthog/hogql/printer.py:0: error: Non-overlapping equality check (left operand
 posthog/hogql/printer.py:0: error: Argument 2 to "_get_materialized_column" of "_Printer" has incompatible type "str | int"; expected "str"  [arg-type]
 posthog/hogql/printer.py:0: error: Statement is unreachable  [unreachable]
 posthog/hogql/printer.py:0: error: Argument 1 to "_print_identifier" of "_Printer" has incompatible type "str | None"; expected "str"  [arg-type]
+posthog/api/organization.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
+posthog/api/organization.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
 posthog/hogql/hogql.py:0: error: Item "None" of "JoinExpr | None" has no attribute "alias"  [union-attr]
 posthog/api/team.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
 posthog/models/property/util.py:0: error: Argument 3 to "format_filter_query" has incompatible type "HogQLContext | None"; expected "HogQLContext"  [arg-type]

posthog/hogql/placeholders.py

@@ -20,6 +20,9 @@ def __init__(self):
         super().__init__()
         self.found: set[str] = set()
 
+    def visit_cte(self, node: ast.CTE):
+        super().visit(node.expr)
+
     def visit_placeholder(self, node: ast.Placeholder):
         self.found.add(node.field)
 

posthog/hogql/test/test_placeholders.py

@@ -1,6 +1,7 @@
+from typing import cast
 from posthog.hogql import ast
 from posthog.hogql.errors import HogQLException
-from posthog.hogql.parser import parse_expr
+from posthog.hogql.parser import parse_expr, parse_select
 from posthog.hogql.placeholders import replace_placeholders, find_placeholders
 from posthog.test.base import BaseTest
 
@@ -69,3 +70,16 @@ def test_assert_no_placeholders(self):
             "Placeholders, such as {foo}, are not supported in this context",
             str(context.exception),
         )
+
+    def test_replace_placeholders_with_cte(self):
+        expr = cast(ast.SelectQuery, parse_select("with test as (select {foo}) select * from test"))
+
+        assert expr.ctes is not None and expr.ctes["test"] is not None
+        assert isinstance(expr.ctes["test"].expr, ast.SelectQuery)
+        assert isinstance(expr.ctes["test"].expr.select[0], ast.Placeholder)
+
+        expr2 = cast(ast.SelectQuery, replace_placeholders(expr, {"foo": ast.Constant(value=1)}))
+
+        assert expr2.ctes is not None and expr2.ctes["test"] is not None
+        assert isinstance(expr2.ctes["test"].expr, ast.SelectQuery)
+        assert isinstance(expr2.ctes["test"].expr.select[0], ast.Constant)

posthog/hogql/visitor.py

@@ -115,6 +115,9 @@ def visit_join_expr(self, node: ast.JoinExpr):
     def visit_select_query(self, node: ast.SelectQuery):
         # :TRICKY: when adding new fields, also add them to visit_select_query of resolver.py
         self.visit(node.select_from)
+        if node.ctes is not None:
+            for expr in list(node.ctes.values()):
+                self.visit(expr)
         for expr in node.array_join_list or []:
             self.visit(expr)
         for expr in node.select or []: