From d790ade662523199e3059f2149143908dbb91e19 Mon Sep 17 00:00:00 2001 From: Eric Duong Date: Tue, 2 Apr 2024 12:10:02 -0400 Subject: [PATCH] chore(data-warehouse): joins on views (#21151) * working * delete * typing * types * types * typing * Update query snapshots * ordering * more types * add test * update baseline * Added logic to has_child too * Added a new SelectViewType * Updated mypy and fixed type * Fixed alias * fix resolver --------- Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Tom Owers --- mypy-baseline.txt | 21 ++++--- .../api/test/__snapshots__/test_decide.ambr | 3 + posthog/hogql/ast.py | 57 +++++++++++++++++-- posthog/hogql/autocomplete.py | 4 +- posthog/hogql/database/database.py | 2 +- posthog/hogql/database/models.py | 9 +-- posthog/hogql/database/test/test_database.py | 34 ++++++++++- posthog/hogql/printer.py | 10 +++- posthog/hogql/resolver.py | 19 ++++++- posthog/hogql/resolver_utils.py | 2 + posthog/hogql/visitor.py | 4 ++ 11 files changed, 138 insertions(+), 27 deletions(-) diff --git a/mypy-baseline.txt b/mypy-baseline.txt index b0ce8847d444c..58dc6f0906860 100644 --- a/mypy-baseline.txt +++ b/mypy-baseline.txt @@ -9,6 +9,7 @@ posthog/temporal/data_imports/pipelines/zendesk/helpers.py:0: error: Argument 1 posthog/temporal/data_imports/pipelines/zendesk/helpers.py:0: error: Argument 1 to "ensure_pendulum_datetime" has incompatible type "DateTime | Date | datetime | date | str | float | int | None"; expected "DateTime | Date | datetime | date | str | float | int" [arg-type] posthog/temporal/data_imports/pipelines/zendesk/helpers.py:0: error: Item "None" of "DateTime | None" has no attribute "int_timestamp" [union-attr] posthog/temporal/data_imports/pipelines/zendesk/helpers.py:0: error: Argument 1 to "ensure_pendulum_datetime" has incompatible type "str | None"; expected "DateTime | Date | datetime | date | str | float | int" [arg-type] +posthog/hogql/modifiers.py:0: error: Incompatible types in assignment (expression has type "PersonOnEventsMode", variable has type "PersonsOnEventsMode | None") [assignment] posthog/hogql/database/argmax.py:0: error: Argument "chain" to "Field" has incompatible type "list[str]"; expected "list[str | int]" [arg-type] posthog/hogql/database/argmax.py:0: note: "List" is invariant -- see https://mypy.readthedocs.io/en/stable/common_issues.html#variance posthog/hogql/database/argmax.py:0: note: Consider using "Sequence" instead, which is covariant @@ -16,9 +17,6 @@ posthog/hogql/database/argmax.py:0: error: Unsupported operand types for + ("lis 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: Argument "chain" to "FieldTraverserType" has incompatible type "list[str]"; expected "list[str | int]" [arg-type] -posthog/hogql/ast.py:0: note: "List" is invariant -- see https://mypy.readthedocs.io/en/stable/common_issues.html#variance -posthog/hogql/ast.py:0: note: Consider using "Sequence" instead, which is covariant posthog/hogql/ast.py:0: error: Incompatible return value type (got "bool | None", expected "bool") [return-value] posthog/hogql/visitor.py:0: error: Statement is unreachable [unreachable] posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "Type | None"; expected "AST" [arg-type] @@ -41,8 +39,8 @@ posthog/hogql/visitor.py:0: error: Incompatible types in assignment (expression 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: 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] posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "WindowFrameExpr | None"; expected "AST" [arg-type] posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "WindowFrameExpr | None"; expected "AST" [arg-type] posthog/hogql/visitor.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "WindowExpr | None"; expected "AST" [arg-type] @@ -148,7 +146,6 @@ posthog/hogql_queries/legacy_compatibility/filter_to_query.py:0: error: Dict ent posthog/hogql_queries/legacy_compatibility/filter_to_query.py:0: error: Dict entry 0 has incompatible type "str": "LifecycleFilter"; expected "str": "TrendsFilter" [dict-item] posthog/hogql_queries/legacy_compatibility/filter_to_query.py:0: error: Dict entry 0 has incompatible type "str": "StickinessFilter"; expected "str": "TrendsFilter" [dict-item] posthog/hogql_queries/legacy_compatibility/feature_flag.py:0: error: Item "AnonymousUser" of "User | AnonymousUser" has no attribute "email" [union-attr] -posthog/hogql/modifiers.py:0: error: Incompatible types in assignment (expression has type "PersonOnEventsMode", variable has type "PersonsOnEventsMode | None") [assignment] posthog/hogql/functions/cohort.py:0: error: Argument 1 to "escape_clickhouse_string" has incompatible type "str | None"; expected "float | int | str | list[Any] | tuple[Any, ...] | date | datetime | UUID | UUIDT" [arg-type] posthog/hogql/functions/cohort.py:0: error: Argument 1 to "escape_clickhouse_string" has incompatible type "str | None"; expected "float | int | str | list[Any] | tuple[Any, ...] | date | datetime | UUID | UUIDT" [arg-type] posthog/hogql/functions/cohort.py:0: error: Incompatible types in assignment (expression has type "ValuesQuerySet[Cohort, tuple[int, bool | None]]", variable has type "ValuesQuerySet[Cohort, tuple[int, bool | None, str | None]]") [assignment] @@ -215,7 +212,7 @@ posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression posthog/hogql/resolver.py:0: error: Argument "alias" to "TableAliasType" has incompatible type "str | int"; expected "str" [arg-type] posthog/hogql/resolver.py:0: error: Argument "table_type" to "TableAliasType" has incompatible type "LazyTableType"; expected "TableType" [arg-type] posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "LazyTableType", variable has type "TableAliasType") [assignment] -posthog/hogql/resolver.py:0: error: Invalid index type "str | int" for "dict[str, BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType]"; expected type "str" [index] +posthog/hogql/resolver.py:0: error: Invalid index type "str | int" for "dict[str, BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType | SelectViewType]"; expected type "str" [index] posthog/hogql/resolver.py:0: error: Argument 1 to "clone_expr" has incompatible type "SelectQuery | SelectUnionQuery | Field | None"; expected "Expr" [arg-type] posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "JoinExpr | None") [assignment] posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has incompatible type "JoinExpr | None"; expected "Expr" [arg-type] @@ -226,11 +223,12 @@ posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has inco posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SampleExpr | None") [assignment] posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has incompatible type "SampleExpr | None"; expected "Expr" [arg-type] posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "SelectQuery | SelectUnionQuery | Field | None"; expected "AST" [arg-type] +posthog/hogql/resolver.py:0: error: Argument "select_query_type" to "SelectViewType" has incompatible type "SelectQueryType | None"; expected "SelectQueryType | SelectUnionQueryType" [arg-type] posthog/hogql/resolver.py:0: error: Item "None" of "SelectQuery | SelectUnionQuery | Field | None" has no attribute "type" [union-attr] posthog/hogql/resolver.py:0: error: Argument "select_query_type" to "SelectQueryAliasType" has incompatible type "Type | Any | None"; expected "SelectQueryType | SelectUnionQueryType" [arg-type] posthog/hogql/resolver.py:0: error: Item "None" of "SelectQuery | SelectUnionQuery | Field | None" has no attribute "type" [union-attr] -posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "Type | Any | None", variable has type "BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType | None") [assignment] -posthog/hogql/resolver.py:0: error: Argument 1 to "append" of "list" has incompatible type "BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType | None"; expected "SelectQueryType | SelectUnionQueryType" [arg-type] +posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "Type | Any | None", variable has type "BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType | SelectViewType | None") [assignment] +posthog/hogql/resolver.py:0: error: Argument 1 to "append" of "list" has incompatible type "BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType | SelectViewType | None"; expected "SelectQueryType | SelectUnionQueryType" [arg-type] posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "JoinExpr | None") [assignment] posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has incompatible type "JoinExpr | None"; expected "Expr" [arg-type] posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "JoinConstraint | None") [assignment] @@ -238,7 +236,7 @@ posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has inco posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SampleExpr | None") [assignment] posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has incompatible type "SampleExpr | None"; expected "Expr" [arg-type] posthog/hogql/resolver.py:0: error: Argument 2 to "convert_hogqlx_tag" has incompatible type "int | None"; expected "int" [arg-type] -posthog/hogql/resolver.py:0: error: Invalid index type "str | int" for "dict[str, BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType]"; expected type "str" [index] +posthog/hogql/resolver.py:0: error: Invalid index type "str | int" for "dict[str, BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType | SelectViewType]"; expected type "str" [index] posthog/hogql/resolver.py:0: error: Argument 2 to "lookup_field_by_name" has incompatible type "str | int"; expected "str" [arg-type] posthog/hogql/resolver.py:0: error: Argument 2 to "lookup_cte_by_name" has incompatible type "str | int"; expected "str" [arg-type] posthog/hogql/resolver.py:0: error: Argument 1 to "get_child" of "Type" has incompatible type "str | int"; expected "str" [arg-type] @@ -260,7 +258,7 @@ posthog/hogql/transforms/lazy_tables.py:0: error: Non-overlapping equality check posthog/hogql/transforms/lazy_tables.py:0: error: Name "chain" already defined on line 0 [no-redef] posthog/hogql/transforms/lazy_tables.py:0: error: Subclass of "TableType" and "LazyTableType" cannot exist: would have incompatible method signatures [unreachable] posthog/hogql/transforms/lazy_tables.py:0: error: Statement is unreachable [unreachable] -posthog/hogql/transforms/lazy_tables.py:0: error: Incompatible types in assignment (expression has type "BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType", variable has type "SelectQueryAliasType | None") [assignment] +posthog/hogql/transforms/lazy_tables.py:0: error: Incompatible types in assignment (expression has type "BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType | SelectViewType", variable has type "SelectQueryAliasType | None") [assignment] posthog/hogql/transforms/in_cohort.py:0: error: Incompatible default for argument "context" (default has type "None", argument has type "HogQLContext") [assignment] posthog/hogql/transforms/in_cohort.py:0: note: PEP 484 prohibits implicit Optional. Accordingly, mypy has changed its default to no_implicit_optional=True posthog/hogql/transforms/in_cohort.py:0: note: Use https://github.com/hauntsaninja/no_implicit_optional to automatically upgrade your codebase @@ -294,6 +292,7 @@ posthog/hogql/printer.py:0: error: Subclass of "TableType" and "LazyTableType" c posthog/hogql/printer.py:0: error: Argument 1 to "visit" of "_Printer" has incompatible type "SelectQuery | SelectUnionQuery | Field | None"; expected "AST" [arg-type] posthog/hogql/printer.py:0: error: Argument 1 to "visit" of "_Printer" has incompatible type "SelectQuery | SelectUnionQuery | Field | None"; expected "AST" [arg-type] posthog/hogql/printer.py:0: error: Argument 1 to "visit" of "_Printer" has incompatible type "SelectQuery | SelectUnionQuery | Field | None"; expected "AST" [arg-type] +posthog/hogql/printer.py:0: error: Argument 1 to "visit" of "_Printer" has incompatible type "SelectQuery | SelectUnionQuery | Field | None"; expected "AST" [arg-type] posthog/hogql/printer.py:0: error: Argument 1 to "_print_escaped_string" of "_Printer" has incompatible type "int | float | UUID | date | None"; expected "float | int | str | list[Any] | tuple[Any, ...] | datetime | date" [arg-type] posthog/hogql/printer.py:0: error: Argument 1 to "join" of "str" has incompatible type "list[str | int]"; expected "Iterable[str]" [arg-type] posthog/hogql/printer.py:0: error: Name "args" already defined on line 0 [no-redef] diff --git a/posthog/api/test/__snapshots__/test_decide.ambr b/posthog/api/test/__snapshots__/test_decide.ambr index 1bb6ffa074d20..7636905bca90b 100644 --- a/posthog/api/test/__snapshots__/test_decide.ambr +++ b/posthog/api/test/__snapshots__/test_decide.ambr @@ -1,4 +1,7 @@ # serializer version: 1 +# name: TestDatabaseCheckForDecide.test_decide_doesnt_error_out_when_database_is_down_and_database_check_isnt_cached + 'SELECT 1' +# --- # name: TestDecide.test_decide_doesnt_error_out_when_database_is_down ''' SELECT "posthog_user"."id", diff --git a/posthog/hogql/ast.py b/posthog/hogql/ast.py index 806226b8f1b9e..9b5daef7a8d9b 100644 --- a/posthog/hogql/ast.py +++ b/posthog/hogql/ast.py @@ -85,6 +85,11 @@ def get_child(self, name: str, context: HogQLContext) -> Type: raise HogQLException(f"Field not found: {name}") +TableOrSelectType = Union[ + BaseTableType, "SelectUnionQueryType", "SelectQueryType", "SelectQueryAliasType", "SelectViewType" +] + + @dataclass(kw_only=True) class TableType(BaseTableType): table: Table @@ -104,7 +109,7 @@ def resolve_database_table(self, context: HogQLContext) -> Table: @dataclass(kw_only=True) class LazyJoinType(BaseTableType): - table_type: BaseTableType + table_type: TableOrSelectType field: str lazy_join: LazyJoin @@ -122,7 +127,7 @@ def resolve_database_table(self, context: HogQLContext) -> Table: @dataclass(kw_only=True) class VirtualTableType(BaseTableType): - table_type: BaseTableType + table_type: TableOrSelectType field: str virtual_table: VirtualTable @@ -133,9 +138,6 @@ def has_child(self, name: str, context: HogQLContext) -> bool: return self.virtual_table.has_field(name) -TableOrSelectType = Union[BaseTableType, "SelectUnionQueryType", "SelectQueryType", "SelectQueryAliasType"] - - @dataclass(kw_only=True) class SelectQueryType(Type): """Type and new enclosed scope for a select query. Contains information about all tables and columns in the query.""" @@ -183,6 +185,49 @@ def has_child(self, name: str, context: HogQLContext) -> bool: return self.types[0].has_child(name, context) +@dataclass(kw_only=True) +class SelectViewType(Type): + view_name: str + alias: str + select_query_type: SelectQueryType | SelectUnionQueryType + + def get_child(self, name: str, context: HogQLContext) -> Type: + if name == "*": + return AsteriskType(table_type=self) + if self.select_query_type.has_child(name, context): + return FieldType(name=name, table_type=self) + if self.view_name: + if context.database is None: + raise HogQLException("Database must be set for queries with views") + + field = context.database.get_table(self.view_name).get_field(name) + + if isinstance(field, LazyJoin): + return LazyJoinType(table_type=self, field=name, lazy_join=field) + if isinstance(field, LazyTable): + return LazyTableType(table=field) + if isinstance(field, FieldTraverser): + return FieldTraverserType(table_type=self, chain=field.chain) + if isinstance(field, VirtualTable): + return VirtualTableType(table_type=self, field=name, virtual_table=field) + if isinstance(field, ExpressionField): + return ExpressionFieldType(table_type=self, name=name, expr=field.expr) + return FieldType(name=name, table_type=self) + raise HogQLException(f"Field {name} not found on view query with name {self.view_name}") + + def has_child(self, name: str, context: HogQLContext) -> bool: + if self.view_name: + if context.database is None: + raise HogQLException("Database must be set for queries with views") + try: + context.database.get_table(self.view_name).get_field(name) + return True + except Exception: + pass + + return self.select_query_type.has_child(name, context) + + @dataclass(kw_only=True) class SelectQueryAliasType(Type): alias: str @@ -193,6 +238,7 @@ def get_child(self, name: str, context: HogQLContext) -> Type: return AsteriskType(table_type=self) if self.select_query_type.has_child(name, context): return FieldType(name=name, table_type=self) + raise HogQLException(f"Field {name} not found on query with alias {self.alias}") def has_child(self, name: str, context: HogQLContext) -> bool: @@ -575,6 +621,7 @@ class SelectQuery(Expr): limit_with_ties: Optional[bool] = None offset: Optional[Expr] = None settings: Optional[HogQLQuerySettings] = None + view_name: Optional[str] = None @dataclass(kw_only=True) diff --git a/posthog/hogql/autocomplete.py b/posthog/hogql/autocomplete.py index a7339a80fafd5..b6d003c1ac88d 100644 --- a/posthog/hogql/autocomplete.py +++ b/posthog/hogql/autocomplete.py @@ -216,9 +216,9 @@ def resolve_table_field_traversers(table: Table, context: HogQLContext) -> Table current_table_or_field: FieldOrTable = new_table for chain in field.chain: if isinstance(current_table_or_field, Table): - chain_field = current_table_or_field.fields.get(chain) + chain_field = current_table_or_field.fields.get(str(chain)) elif isinstance(current_table_or_field, LazyJoin): - chain_field = current_table_or_field.resolve_table(context).fields.get(chain) + chain_field = current_table_or_field.resolve_table(context).fields.get(str(chain)) elif isinstance(current_table_or_field, DatabaseField): chain_field = current_table_or_field else: diff --git a/posthog/hogql/database/database.py b/posthog/hogql/database/database.py index afeac3c26a143..df536b6c9ae2d 100644 --- a/posthog/hogql/database/database.py +++ b/posthog/hogql/database/database.py @@ -352,7 +352,7 @@ class _SerializedFieldBase(TypedDict): class SerializedField(_SerializedFieldBase, total=False): fields: List[str] table: str - chain: List[str] + chain: List[str | int] def serialize_database(context: HogQLContext) -> Dict[str, List[SerializedField]]: diff --git a/posthog/hogql/database/models.py b/posthog/hogql/database/models.py index d2da7868a7f9c..46e0ba2129791 100644 --- a/posthog/hogql/database/models.py +++ b/posthog/hogql/database/models.py @@ -65,17 +65,18 @@ class ExpressionField(DatabaseField): class FieldTraverser(FieldOrTable): model_config = ConfigDict(extra="forbid") - chain: List[str] + chain: List[str | int] class Table(FieldOrTable): fields: Dict[str, FieldOrTable] model_config = ConfigDict(extra="forbid") - def has_field(self, name: str) -> bool: - return name in self.fields + def has_field(self, name: str | int) -> bool: + return str(name) in self.fields - def get_field(self, name: str) -> FieldOrTable: + def get_field(self, name: str | int) -> FieldOrTable: + name = str(name) if self.has_field(name): return self.fields[name] raise Exception(f'Field "{name}" not found on table {self.__class__.__name__}') diff --git a/posthog/hogql/database/test/test_database.py b/posthog/hogql/database/test/test_database.py index 99810ca54e58a..115f7282661b5 100644 --- a/posthog/hogql/database/test/test_database.py +++ b/posthog/hogql/database/test/test_database.py @@ -17,7 +17,7 @@ from posthog.models.organization import Organization from posthog.models.team.team import Team from posthog.test.base import BaseTest -from posthog.warehouse.models import DataWarehouseTable, DataWarehouseCredential +from posthog.warehouse.models import DataWarehouseTable, DataWarehouseCredential, DataWarehouseSavedQuery from posthog.hogql.query import execute_hogql_query from posthog.warehouse.models.join import DataWarehouseJoin @@ -288,3 +288,35 @@ def test_database_warehouse_joins_persons_poe_v2(self): assert poe.fields["some_field"] is not None print_ast(parse_select("select person.some_field.key from events"), context, dialect="clickhouse") + + def test_database_warehouse_joins_on_view(self): + DataWarehouseSavedQuery.objects.create( + team=self.team, + name="event_view", + query={"query": "SELECT event AS event from events"}, + columns={"event": "String"}, + ) + DataWarehouseJoin.objects.create( + team=self.team, + source_table_name="event_view", + source_table_key="event", + joining_table_name="groups", + joining_table_key="key", + field_name="some_field", + ) + + db = create_hogql_database(team_id=self.team.pk) + context = HogQLContext( + team_id=self.team.pk, + enable_select_queries=True, + database=db, + ) + + sql = "select event_view.some_field.key from event_view" + print_ast(parse_select(sql), context, dialect="clickhouse") + + sql = "select some_field.key from event_view" + print_ast(parse_select(sql), context, dialect="clickhouse") + + sql = "select e.some_field.key from event_view as e" + print_ast(parse_select(sql), context, dialect="clickhouse") diff --git a/posthog/hogql/printer.py b/posthog/hogql/printer.py index 90cbc6663519e..c3444540e9627 100644 --- a/posthog/hogql/printer.py +++ b/posthog/hogql/printer.py @@ -424,6 +424,10 @@ def visit_join_expr(self, node: ast.JoinExpr) -> JoinExprResponse: elif isinstance(node.type, ast.SelectUnionQueryType): join_strings.append(self.visit(node.table)) + elif isinstance(node.type, ast.SelectViewType) and node.alias is not None: + join_strings.append(self.visit(node.table)) + join_strings.append(f"AS {self._print_identifier(node.alias)}") + elif isinstance(node.type, ast.SelectQueryAliasType) and node.alias is not None: join_strings.append(self.visit(node.table)) join_strings.append(f"AS {self._print_identifier(node.alias)}") @@ -968,10 +972,11 @@ def visit_field_type(self, type: ast.FieldType): elif ( isinstance(type.table_type, ast.SelectQueryType) or isinstance(type.table_type, ast.SelectQueryAliasType) + or isinstance(type.table_type, ast.SelectViewType) or isinstance(type.table_type, ast.SelectUnionQueryType) ): field_sql = self._print_identifier(type.name) - if isinstance(type.table_type, ast.SelectQueryAliasType): + if isinstance(type.table_type, ast.SelectQueryAliasType) or isinstance(type.table_type, ast.SelectViewType): field_sql = f"{self.visit(type.table_type)}.{field_sql}" # :KLUDGE: Legacy person properties handling. Only used within non-HogQL queries, such as insights. @@ -1066,6 +1071,9 @@ def visit_ratio_expr(self, node: ast.RatioExpr): def visit_select_query_alias_type(self, type: ast.SelectQueryAliasType): return self._print_identifier(type.alias) + def visit_select_view_type(self, type: ast.SelectViewType): + return self._print_identifier(type.alias) + def visit_field_alias_type(self, type: ast.FieldAliasType): return self._print_identifier(type.alias) diff --git a/posthog/hogql/resolver.py b/posthog/hogql/resolver.py index 19bbddb9f65b7..cebd0e95e8fc0 100644 --- a/posthog/hogql/resolver.py +++ b/posthog/hogql/resolver.py @@ -207,6 +207,7 @@ def visit_select_query(self, node: ast.SelectQuery): {name: self.visit(expr) for name, expr in node.window_exprs.items()} if node.window_exprs else None ) new_node.settings = node.settings.model_copy() if node.settings is not None else None + new_node.view_name = node.view_name self.scopes.pop() @@ -222,9 +223,10 @@ def _asterisk_columns(self, asterisk: ast.AsteriskType) -> List[ast.Expr]: isinstance(asterisk.table_type, ast.SelectUnionQueryType) or isinstance(asterisk.table_type, ast.SelectQueryType) or isinstance(asterisk.table_type, ast.SelectQueryAliasType) + or isinstance(asterisk.table_type, ast.SelectViewType) ): select = asterisk.table_type - while isinstance(select, ast.SelectQueryAliasType): + while isinstance(select, ast.SelectQueryAliasType) or isinstance(select, ast.SelectViewType): select = select.select_query_type if isinstance(select, ast.SelectUnionQueryType): select = select.types[0] @@ -276,6 +278,10 @@ def visit_join_expr(self, node: ast.JoinExpr): raise ResolverException("Nested views are not supported") node.table = parse_select(str(database_table.query)) + + if isinstance(node.table, ast.SelectQuery): + node.table.view_name = database_table.name + node.alias = table_alias or database_table.name node = self.visit(node) @@ -328,7 +334,16 @@ def visit_join_expr(self, node: ast.JoinExpr): node = cast(ast.JoinExpr, clone_expr(node)) node.table = super().visit(node.table) - if node.alias is not None: + if isinstance(node.table, ast.SelectQuery) and node.table.view_name is not None and node.alias is not None: + if node.alias in scope.tables: + raise ResolverException( + f'Already have joined a table called "{node.alias}". Can\'t join another one with the same name.' + ) + node.type = ast.SelectViewType( + alias=node.alias, view_name=node.table.view_name, select_query_type=node.table.type + ) + scope.tables[node.alias] = node.type + elif node.alias is not None: if node.alias in scope.tables: raise ResolverException( f'Already have joined a table called "{node.alias}". Can\'t join another one with the same name.' diff --git a/posthog/hogql/resolver_utils.py b/posthog/hogql/resolver_utils.py index 7d39fbb59e36d..457913e7006f3 100644 --- a/posthog/hogql/resolver_utils.py +++ b/posthog/hogql/resolver_utils.py @@ -43,6 +43,8 @@ def get_long_table_name(select: ast.SelectQueryType, type: ast.Type) -> str: return type.alias elif isinstance(type, ast.SelectQueryAliasType): return type.alias + elif isinstance(type, ast.SelectViewType): + return type.alias elif isinstance(type, ast.LazyJoinType): return f"{get_long_table_name(select, type.table_type)}__{type.field}" elif isinstance(type, ast.VirtualTableType): diff --git a/posthog/hogql/visitor.py b/posthog/hogql/visitor.py index 2bf968abf2ab0..11580e9ec136e 100644 --- a/posthog/hogql/visitor.py +++ b/posthog/hogql/visitor.py @@ -187,6 +187,9 @@ def visit_table_alias_type(self, node: ast.TableAliasType): def visit_select_query_alias_type(self, node: ast.SelectQueryAliasType): self.visit(node.select_query_type) + def visit_select_view_type(self, node: ast.SelectViewType): + self.visit(node.select_query_type) + def visit_asterisk_type(self, node: ast.AsteriskType): self.visit(node.table_type) @@ -485,6 +488,7 @@ def visit_select_query(self, node: ast.SelectQuery): if node.window_exprs else None, settings=node.settings.model_copy() if node.settings is not None else None, + view_name=node.view_name, ) def visit_select_union_query(self, node: ast.SelectUnionQuery):