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 <[email protected]>
PostHog · Apr 2, 2024 · d790ade · d790ade
1 parent edb3a46
commit d790ade
Show file tree

Hide file tree

Showing 11 changed files with 138 additions and 27 deletions.
diff --git a/mypy-baseline.txt b/mypy-baseline.txt
@@ -9,16 +9,14 @@ 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
 posthog/hogql/database/argmax.py:0: error: Unsupported operand types for + ("list[str]" and "list[str | int]")  [operator]
 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,19 +223,20 @@ 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]
 posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has incompatible type "JoinConstraint | None"; expected "Expr"  [arg-type]
 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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__}')