feat(hog): print ast in console (#25608)

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]

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):

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

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")
 

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)
 
 

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)")

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,

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,

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: