From 0adc0490ed634b28818651ffa1958c955a0c55fa Mon Sep 17 00:00:00 2001 From: Tom Owers Date: Mon, 12 Feb 2024 17:04:25 +0000 Subject: [PATCH] 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 | 10 ++++++++-- posthog/hogql/placeholders.py | 3 +++ posthog/hogql/test/test_placeholders.py | 16 +++++++++++++++- posthog/hogql/visitor.py | 3 +++ 4 files changed, 29 insertions(+), 3 deletions(-) diff --git a/mypy-baseline.txt b/mypy-baseline.txt index 8b56066703dcd..5976a1cc73bd2 100644 --- a/mypy-baseline.txt +++ b/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] diff --git a/posthog/hogql/placeholders.py b/posthog/hogql/placeholders.py index bd63ce32754c0..ec0a75f04e5bb 100644 --- a/posthog/hogql/placeholders.py +++ b/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) diff --git a/posthog/hogql/test/test_placeholders.py b/posthog/hogql/test/test_placeholders.py index 88c92ebfc8fe8..752400776bcc4 100644 --- a/posthog/hogql/test/test_placeholders.py +++ b/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) diff --git a/posthog/hogql/visitor.py b/posthog/hogql/visitor.py index 5eb93400ce555..9acb6fec87db7 100644 --- a/posthog/hogql/visitor.py +++ b/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 []: