Skip to content

Commit

Permalink
chore(data-warehouse): joins on views (#21151)
Browse files Browse the repository at this point in the history
* 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 <owerstom@gmail.com>
  • Loading branch information
3 people authored Apr 2, 2024
1 parent edb3a46 commit d790ade
Show file tree
Hide file tree
Showing 11 changed files with 138 additions and 27 deletions.
21 changes: 10 additions & 11 deletions mypy-baseline.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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]
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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]
Expand All @@ -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]
Expand All @@ -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
Expand Down Expand Up @@ -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]
Expand Down
3 changes: 3 additions & 0 deletions posthog/api/test/__snapshots__/test_decide.ambr
Original file line number Diff line number Diff line change
@@ -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",
Expand Down
57 changes: 52 additions & 5 deletions posthog/hogql/ast.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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

Expand All @@ -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."""
Expand Down Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
4 changes: 2 additions & 2 deletions posthog/hogql/autocomplete.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion posthog/hogql/database/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]]:
Expand Down
9 changes: 5 additions & 4 deletions posthog/hogql/database/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__}')
Expand Down
Loading

0 comments on commit d790ade

Please sign in to comment.