From e8031f7798b379289eec347e7897da2469d83177 Mon Sep 17 00:00:00 2001 From: Marius Andra Date: Mon, 21 Oct 2024 09:28:06 +0200 Subject: [PATCH] feat(hog): print ast in console (#25608) --- mypy-baseline.txt | 25 +++++++++++----------- posthog/hogql/base.py | 20 ++++++++++++++++- posthog/hogql/context.py | 2 +- posthog/hogql/printer.py | 16 ++++++++------ posthog/hogql/resolver.py | 5 +++-- posthog/hogql/test/test_printer.py | 5 +++++ posthog/hogql/transforms/in_cohort.py | 5 +++-- posthog/hogql/transforms/lazy_tables.py | 3 ++- posthog/hogql/transforms/property_types.py | 2 +- 9 files changed, 56 insertions(+), 27 deletions(-) diff --git a/mypy-baseline.txt b/mypy-baseline.txt index f57273827b467..ae7bf250f3f0a 100644 --- a/mypy-baseline.txt +++ b/mypy-baseline.txt @@ -61,10 +61,6 @@ posthog/settings/data_stores.py:0: error: Name "DATABASE_URL" already defined on posthog/plugins/utils.py:0: error: Subclass of "str" and "bytes" cannot exist: would have incompatible method signatures [unreachable] posthog/plugins/utils.py:0: error: Statement is unreachable [unreachable] posthog/models/dashboard.py:0: error: Need type annotation for "insights" [var-annotated] -posthog/hogql/database/schema/numbers.py:0: error: Incompatible types in assignment (expression has type "dict[str, IntegerDatabaseField]", variable has type "dict[str, FieldOrTable]") [assignment] -posthog/hogql/database/schema/numbers.py:0: note: "Dict" is invariant -- see https://mypy.readthedocs.io/en/stable/common_issues.html#variance -posthog/hogql/database/schema/numbers.py:0: note: Consider using "Mapping" instead, which is covariant in the value type -posthog/hogql/ast.py:0: error: Incompatible return value type (got "bool | None", expected "bool") [return-value] posthog/warehouse/data_load/service.py:0: error: Unsupported operand types for >= ("timedelta" and "None") [operator] posthog/warehouse/data_load/service.py:0: note: Left operand is of type "timedelta | None" posthog/warehouse/data_load/service.py:0: error: Incompatible return value type (got "tuple[timedelta | None, timedelta]", expected "tuple[timedelta, timedelta]") [return-value] @@ -77,6 +73,15 @@ posthog/models/subscription.py:0: error: Argument 2 to "SubscriptionResourceInfo posthog/models/exported_asset.py:0: error: Value of type variable "_StrOrPromiseT" of "slugify" cannot be "str | None" [type-var] posthog/models/action/action.py:0: error: Need type annotation for "events" [var-annotated] posthog/models/action/action.py:0: error: Argument 1 to "len" has incompatible type "str | None"; expected "Sized" [arg-type] +posthog/hogql/database/schema/numbers.py:0: error: Incompatible types in assignment (expression has type "dict[str, IntegerDatabaseField]", variable has type "dict[str, FieldOrTable]") [assignment] +posthog/hogql/database/schema/numbers.py:0: note: "Dict" is invariant -- see https://mypy.readthedocs.io/en/stable/common_issues.html#variance +posthog/hogql/database/schema/numbers.py:0: note: Consider using "Mapping" instead, which is covariant in the value type +posthog/hogql/ast.py:0: error: Incompatible return value type (got "bool | None", expected "bool") [return-value] +ee/models/license.py:0: error: Incompatible return value type (got "_T", expected "License | None") [return-value] +ee/models/license.py:0: error: Cannot use a covariant type variable as a parameter [misc] +ee/models/license.py:0: error: "_T" has no attribute "plan" [attr-defined] +ee/models/license.py:0: error: Incompatible return value type (got "str | bool", expected "bool") [return-value] +ee/models/explicit_team_membership.py:0: error: Incompatible return value type (got "int", expected "Level") [return-value] 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: Incompatible types in assignment (expression has type "Expr", variable has type "CTE") [assignment] @@ -85,18 +90,13 @@ posthog/hogql/visitor.py:0: error: Incompatible types in assignment (expression 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 | SelectViewType") [assignment] posthog/hogql/visitor.py:0: error: Incompatible types in assignment (expression has type "Type", variable has type "BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType | SelectViewType") [assignment] -ee/models/license.py:0: error: Incompatible return value type (got "_T", expected "License | None") [return-value] -ee/models/license.py:0: error: Cannot use a covariant type variable as a parameter [misc] -ee/models/license.py:0: error: "_T" has no attribute "plan" [attr-defined] -ee/models/license.py:0: error: Incompatible return value type (got "str | bool", expected "bool") [return-value] -ee/models/explicit_team_membership.py:0: error: Incompatible return value type (got "int", expected "Level") [return-value] -posthog/hogql/resolver_utils.py:0: error: Argument 1 to "lookup_field_by_name" has incompatible type "SelectQueryType | SelectUnionQueryType"; expected "SelectQueryType" [arg-type] posthog/models/filters/mixins/simplify.py:0: error: Incompatible type for lookup 'pk': (got "str | int | list[str]", expected "str | int") [misc] +posthog/hogql/resolver_utils.py:0: error: Argument 1 to "lookup_field_by_name" has incompatible type "SelectQueryType | SelectUnionQueryType"; expected "SelectQueryType" [arg-type] +posthog/helpers/dashboard_templates.py:0: error: Incompatible types in assignment (expression has type "str | None", variable has type "str | Combinable") [assignment] posthog/hogql/parser.py:0: error: Item "None" of "list[Expr] | None" has no attribute "__iter__" (not iterable) [union-attr] posthog/hogql/parser.py:0: error: "None" has no attribute "text" [attr-defined] posthog/hogql/parser.py:0: error: "None" has no attribute "text" [attr-defined] posthog/hogql/parser.py:0: error: Statement is unreachable [unreachable] -posthog/helpers/dashboard_templates.py:0: error: Incompatible types in assignment (expression has type "str | None", variable has type "str | Combinable") [assignment] posthog/hogql/functions/cohort.py:0: error: Incompatible type for lookup 'team_id': (got "int | None", expected "str | int") [misc] posthog/hogql/functions/cohort.py:0: error: Incompatible type for lookup 'team_id': (got "int | None", expected "str | int") [misc] posthog/hogql/database/schema/persons_pdi.py:0: error: Incompatible types in assignment (expression has type "Organization | None", variable has type "Organization") [assignment] @@ -106,8 +106,8 @@ posthog/hogql/database/schema/groups.py:0: note: Consider using "Mapping" instea posthog/hogql/database/schema/groups.py:0: error: Incompatible types in assignment (expression has type "dict[str, DatabaseField]", variable has type "dict[str, FieldOrTable]") [assignment] posthog/hogql/database/schema/groups.py:0: note: "Dict" is invariant -- see https://mypy.readthedocs.io/en/stable/common_issues.html#variance posthog/hogql/database/schema/groups.py:0: note: Consider using "Mapping" instead, which is covariant in the value type -posthog/hogql/database/schema/persons.py:0: error: Incompatible types in assignment (expression has type "Organization | None", variable has type "Organization") [assignment] posthog/batch_exports/service.py:0: error: Argument 4 to "backfill_export" has incompatible type "datetime | None"; expected "datetime" [arg-type] +posthog/hogql/database/schema/persons.py:0: error: Incompatible types in assignment (expression has type "Organization | None", variable has type "Organization") [assignment] posthog/models/team/team.py:0: error: Statement is unreachable [unreachable] posthog/models/team/team.py:0: error: Statement is unreachable [unreachable] posthog/models/hog_functions/hog_function.py:0: error: Argument 1 to "get" of "dict" has incompatible type "str | None"; expected "str" [arg-type] @@ -230,7 +230,6 @@ posthog/hogql/transforms/in_cohort.py:0: error: List item 0 has incompatible typ posthog/hogql/database/database.py:0: error: Argument "week_start_day" to "Database" has incompatible type "int | Any | None"; expected "WeekStartDay | None" [arg-type] posthog/hogql/database/database.py:0: error: "FieldOrTable" has no attribute "fields" [attr-defined] posthog/warehouse/models/datawarehouse_saved_query.py:0: error: Argument 1 to "create_hogql_database" has incompatible type "int | None"; expected "int" [arg-type] -posthog/warehouse/models/datawarehouse_saved_query.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SelectQuery | SelectUnionQuery") [assignment] posthog/models/feature_flag/flag_matching.py:0: error: Statement is unreachable [unreachable] posthog/models/feature_flag/flag_matching.py:0: error: Value expression in dictionary comprehension has incompatible type "int"; expected type "Literal[0, 1, 2, 3, 4]" [misc] posthog/models/feature_flag/flag_matching.py:0: error: Value of type variable "_E" of "ExpressionWrapper" cannot be "object" [type-var] diff --git a/posthog/hogql/base.py b/posthog/hogql/base.py index 6f281e3a393ab..6eede3377bebb 100644 --- a/posthog/hogql/base.py +++ b/posthog/hogql/base.py @@ -1,7 +1,7 @@ import re from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Literal, Optional +from typing import TYPE_CHECKING, Literal, Optional, TypeVar from posthog.hogql.constants import ConstantDataType from posthog.hogql.errors import NotImplementedError @@ -34,6 +34,24 @@ def accept(self, visitor): return visitor.visit_unknown(self) raise NotImplementedError(f"{visitor.__class__.__name__} has no method {method_name}") + def to_hogql(self): + from posthog.hogql.printer import print_prepared_ast + from posthog.hogql.context import HogQLContext + + return print_prepared_ast( + node=self, + context=HogQLContext(enable_select_queries=True, limit_top_select=False), + dialect="hogql", + ) + + def __str__(self): + if isinstance(self, Type): + return super().__str__() + return f"sql({self.to_hogql()})" + + +_T_AST = TypeVar("_T_AST", bound=AST) + @dataclass(kw_only=True) class Type(AST): diff --git a/posthog/hogql/context.py b/posthog/hogql/context.py index e47b854f01ac7..615aedc8d2f6c 100644 --- a/posthog/hogql/context.py +++ b/posthog/hogql/context.py @@ -23,7 +23,7 @@ class HogQLContext: """Context given to a HogQL expression printer""" # Team making the queries - team_id: Optional[int] + team_id: Optional[int] = None # Team making the queries - if team is passed in, then the team isn't queried when creating the database team: Optional["Team"] = None # Virtual database we're querying, will be populated from team_id if not present diff --git a/posthog/hogql/printer.py b/posthog/hogql/printer.py index a77791034803c..1ebe4f229dcf8 100644 --- a/posthog/hogql/printer.py +++ b/posthog/hogql/printer.py @@ -8,7 +8,7 @@ from posthog.clickhouse.property_groups import property_groups from posthog.hogql import ast -from posthog.hogql.base import AST +from posthog.hogql.base import AST, _T_AST from posthog.hogql.constants import ( MAX_SELECT_RETURNED_ROWS, HogQLGlobalSettings, @@ -78,7 +78,7 @@ def to_printed_hogql(query: ast.Expr, team: Team, modifiers: Optional[HogQLQuery def print_ast( - node: ast.Expr, + node: _T_AST, context: HogQLContext, dialect: Literal["hogql", "clickhouse"], stack: Optional[list[ast.SelectQuery]] = None, @@ -99,12 +99,12 @@ def print_ast( def prepare_ast_for_printing( - node: ast.Expr, + node: _T_AST, context: HogQLContext, dialect: Literal["hogql", "clickhouse"], stack: Optional[list[ast.SelectQuery]] = None, settings: Optional[HogQLGlobalSettings] = None, -) -> ast.Expr | None: +) -> _T_AST | None: with context.timings.measure("create_hogql_database"): context.database = context.database or create_hogql_database(context.team_id, context.modifiers, context.team) @@ -166,7 +166,7 @@ def prepare_ast_for_printing( def print_prepared_ast( - node: ast.Expr, + node: _T_AST, context: HogQLContext, dialect: Literal["hogql", "clickhouse"], stack: Optional[list[ast.SelectQuery]] = None, @@ -523,6 +523,10 @@ def visit_join_expr(self, node: ast.JoinExpr) -> JoinExprResponse: join_strings.append(self._print_identifier(node.type.table.to_printed_hogql())) else: raise ImpossibleASTError(f"Unexpected LazyTableType for: {node.type.table.to_printed_hogql()}") + + elif self.dialect == "hogql": + join_strings.append(self.visit(node.table)) + else: raise QueryError( f"Only selecting from a table or a subquery is supported. Unexpected type: {node.type.__class__.__name__}" @@ -921,7 +925,7 @@ def visit_constant(self, node: ast.Constant): return self.context.add_value(node.value) def visit_field(self, node: ast.Field): - if node.type is None: + if node.type is None and self.dialect != "hogql": field = ".".join([self._print_hogql_identifier_or_index(identifier) for identifier in node.chain]) raise ImpossibleASTError(f"Field {field} has no type") diff --git a/posthog/hogql/resolver.py b/posthog/hogql/resolver.py index 06f6fbea9dfc7..b263843598f87 100644 --- a/posthog/hogql/resolver.py +++ b/posthog/hogql/resolver.py @@ -4,6 +4,7 @@ from posthog.hogql import ast from posthog.hogql.ast import ConstantType, FieldTraverserType +from posthog.hogql.base import _T_AST from posthog.hogql.context import HogQLContext from posthog.hogql.database.models import ( FunctionCallTable, @@ -77,11 +78,11 @@ def resolve_types_from_table( def resolve_types( - node: ast.Expr | ast.SelectQuery, + node: _T_AST, context: HogQLContext, dialect: Literal["hogql", "clickhouse"], scopes: Optional[list[ast.SelectQueryType]] = None, -) -> ast.Expr: +) -> _T_AST: return Resolver(scopes=scopes, context=context, dialect=dialect).visit(node) diff --git a/posthog/hogql/test/test_printer.py b/posthog/hogql/test/test_printer.py index fd735a8b87ce4..021dee3bcf947 100644 --- a/posthog/hogql/test/test_printer.py +++ b/posthog/hogql/test/test_printer.py @@ -114,6 +114,11 @@ def test_to_printed_hogql(self): repsponse, f"SELECT\n plus(1, 2),\n 3\nFROM\n events\nLIMIT {MAX_SELECT_RETURNED_ROWS}" ) + def test_print_to_string(self): + assert str(parse_select("select 1 + 2, 3 from events")) == "sql(SELECT plus(1, 2), 3 FROM events)" + assert str(parse_expr("1 + 2")) == "sql(plus(1, 2))" + assert str(parse_expr("unknown_field")) == "sql(unknown_field)" + def test_literals(self): self.assertEqual(self._expr("1 + 2"), "plus(1, 2)") self.assertEqual(self._expr("-1 + 2"), "plus(-1, 2)") diff --git a/posthog/hogql/transforms/in_cohort.py b/posthog/hogql/transforms/in_cohort.py index fec1a4d7ccd16..418b097e1d07d 100644 --- a/posthog/hogql/transforms/in_cohort.py +++ b/posthog/hogql/transforms/in_cohort.py @@ -2,6 +2,7 @@ from posthog.hogql import ast +from posthog.hogql.base import _T_AST from posthog.hogql.context import HogQLContext from posthog.hogql.errors import QueryError from posthog.hogql.escape_sql import escape_clickhouse_string @@ -11,7 +12,7 @@ def resolve_in_cohorts( - node: ast.Expr, + node: _T_AST, dialect: Literal["hogql", "clickhouse"], stack: Optional[list[ast.SelectQuery]] = None, context: HogQLContext = None, @@ -20,7 +21,7 @@ def resolve_in_cohorts( def resolve_in_cohorts_conjoined( - node: ast.Expr, + node: ast.AST, dialect: Literal["hogql", "clickhouse"], context: HogQLContext, stack: Optional[list[ast.SelectQuery]] = None, diff --git a/posthog/hogql/transforms/lazy_tables.py b/posthog/hogql/transforms/lazy_tables.py index aa0ccfada96e2..13fffdbb1a41c 100644 --- a/posthog/hogql/transforms/lazy_tables.py +++ b/posthog/hogql/transforms/lazy_tables.py @@ -2,6 +2,7 @@ from typing import Optional, cast, Literal from posthog.hogql import ast +from posthog.hogql.base import _T_AST from posthog.hogql.context import HogQLContext from posthog.hogql.database.models import LazyTableToAdd, LazyJoinToAdd from posthog.hogql.errors import ResolutionError @@ -13,7 +14,7 @@ # This mutates the nodes def resolve_lazy_tables( - node: ast.Expr, + node: _T_AST, dialect: Literal["hogql", "clickhouse"], stack: Optional[list[ast.SelectQuery]], context: HogQLContext, diff --git a/posthog/hogql/transforms/property_types.py b/posthog/hogql/transforms/property_types.py index c1bfcc622ff75..53d7df7ebf355 100644 --- a/posthog/hogql/transforms/property_types.py +++ b/posthog/hogql/transforms/property_types.py @@ -13,7 +13,7 @@ from posthog.hogql.database.s3_table import S3Table -def build_property_swapper(node: ast.Expr, context: HogQLContext) -> None: +def build_property_swapper(node: ast.AST, context: HogQLContext) -> None: from posthog.models import PropertyDefinition if not context or not context.team_id: