feat(autocomplete): Added support for joined tables and using autocom…

…plete on join constr… (#20304)

* Added suppoer for joined tables and using autocomplete on join constraints

* Added a type for Expression fields

* Fixed mypy

* Fixed mypy

* Update query snapshots

---------

Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>

.github/actions/run-backend-tests/action.yml

@@ -145,7 +145,6 @@ runs:
           run: echo "PYTEST_ARGS=--snapshot-update" >> $GITHUB_ENV # We can only update snapshots within the PostHog org
 
         # Tests
-
         - name: Run FOSS tests
           if: ${{ inputs.segment == 'FOSS' }}
           env:

mypy-baseline.txt

@@ -365,11 +365,9 @@ posthog/hogql/query.py:0: error: Argument 1 to "get_default_limit_for_context" h
 posthog/hogql/query.py:0: error: "SelectQuery" has no attribute "select_queries"  [attr-defined]
 posthog/hogql/query.py:0: error: Subclass of "SelectQuery" and "SelectUnionQuery" cannot exist: would have incompatible method signatures  [unreachable]
 posthog/hogql/autocomplete.py:0: error: Unused "type: ignore" comment  [unused-ignore]
+posthog/hogql/autocomplete.py:0: error: Unused "type: ignore" comment  [unused-ignore]
 posthog/hogql_queries/insights/trends/breakdown_values.py:0: error: Item "SelectUnionQuery" of "SelectQuery | SelectUnionQuery" has no attribute "select"  [union-attr]
 posthog/hogql_queries/insights/trends/breakdown_values.py:0: error: Value of type "list[Any] | None" is not indexable  [index]
-posthog/hogql_queries/insights/funnels/base.py:0: error: Incompatible types in assignment (expression has type "FunnelExclusionEventsNode | FunnelExclusionActionsNode", variable has type "EventsNode | ActionsNode")  [assignment]
-posthog/hogql_queries/insights/funnels/base.py:0: error: Item "EventsNode" of "EventsNode | ActionsNode" has no attribute "funnelFromStep"  [union-attr]
-posthog/hogql_queries/insights/funnels/base.py:0: error: Item "ActionsNode" of "EventsNode | ActionsNode" has no attribute "funnelFromStep"  [union-attr]
 posthog/hogql_queries/sessions_timeline_query_runner.py:0: error: Statement is unreachable  [unreachable]
 posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown_type"  [union-attr]
 posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown_histogram_bin_count"  [union-attr]
@@ -404,6 +402,9 @@ posthog/hogql_queries/events_query_runner.py:0: error: Statement is unreachable
 posthog/hogql/metadata.py:0: error: Argument "metadata_source" to "translate_hogql" has incompatible type "SelectQuery | SelectUnionQuery"; expected "SelectQuery | None"  [arg-type]
 posthog/hogql/metadata.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SelectQuery | SelectUnionQuery")  [assignment]
 posthog/queries/breakdown_props.py:0: error: Argument 1 to "translate_hogql" has incompatible type "str | int"; expected "str"  [arg-type]
+posthog/hogql_queries/insights/funnels/base.py:0: error: Incompatible types in assignment (expression has type "FunnelExclusionEventsNode | FunnelExclusionActionsNode", variable has type "EventsNode | ActionsNode")  [assignment]
+posthog/hogql_queries/insights/funnels/base.py:0: error: Item "EventsNode" of "EventsNode | ActionsNode" has no attribute "funnelFromStep"  [union-attr]
+posthog/hogql_queries/insights/funnels/base.py:0: error: Item "ActionsNode" of "EventsNode | ActionsNode" has no attribute "funnelFromStep"  [union-attr]
 posthog/queries/funnels/base.py:0: error: "HogQLContext" has no attribute "person_on_events_mode"  [attr-defined]
 posthog/queries/funnels/base.py:0: error: Argument 1 to "translate_hogql" has incompatible type "str | int"; expected "str"  [arg-type]
 ee/clickhouse/queries/funnels/funnel_correlation.py:0: error: Statement is unreachable  [unreachable]

posthog/api/test/dashboards/test_dashboard.py

@@ -905,6 +905,7 @@ def test_dashboard_duplication(self):
         existing_dashboard_item_id_set = {tile1.pk, tile2.pk}
         response_item_id_set = set(map(lambda x: x.get("id", None), response["tiles"]))
         # check both sets are disjoint to verify that the new items' ids are different than the existing items
+
         self.assertTrue(existing_dashboard_item_id_set.isdisjoint(response_item_id_set))
 
         for item in response["tiles"]:

posthog/hogql/autocomplete.py

@@ -97,6 +97,8 @@ def convert_field_or_table_to_type_string(field_or_table: FieldOrTable) -> str |
         return "Date"
     if isinstance(field_or_table, ast.StringJSONDatabaseField):
         return "Object"
+    if isinstance(field_or_table, ast.ExpressionField):
+        return "Expression"
     if isinstance(field_or_table, (ast.Table, ast.LazyJoin)):
         return "Table"
 
@@ -181,6 +183,26 @@ def to_printed_hogql(self):
     return None
 
 
+def get_tables_aliases(query: ast.SelectQuery, context: HogQLContext) -> Dict[str, ast.Table]:
+    tables: Dict[str, ast.Table] = {}
+
+    if query.select_from is not None and query.select_from.alias is not None:
+        table = get_table(context, query.select_from, query.ctes)
+        if table is not None:
+            tables[query.select_from.alias] = table
+
+    if query.select_from is not None and query.select_from.next_join is not None:
+        next_join: ast.JoinExpr | None = query.select_from.next_join
+        while next_join is not None:
+            if next_join.alias is not None:
+                table = get_table(context, next_join, query.ctes)
+                if table is not None:
+                    tables[next_join.alias] = table
+            next_join = next_join.next_join
+
+    return tables
+
+
 # Replaces all ast.FieldTraverser with the underlying node
 def resolve_table_field_traversers(table: Table) -> Table:
     new_table = deepcopy(table)
@@ -316,20 +338,22 @@ def get_hogql_autocomplete(query: HogQLAutocomplete, team: Team) -> HogQLAutocom
                 chain_len = len(node.chain)
                 last_table: Table = table
                 for index, chain_part in enumerate(node.chain):
-                    # TODO: Include joined table aliases
                     # Return just the table alias
                     if table_has_alias and index == 0 and chain_len == 1:
-                        alias = nearest_select.select_from.alias
-                        assert alias is not None
+                        table_aliases = list(get_tables_aliases(nearest_select, context).keys())
                         extend_responses(
-                            keys=[alias],
+                            keys=table_aliases,
                             suggestions=response.suggestions,
                             kind=Kind.Folder,
-                            details=["Table"],
+                            details=["Table"] * len(table_aliases),  # type: ignore
                         )
                         break
 
                     if table_has_alias and index == 0:
+                        tables = get_tables_aliases(nearest_select, context)
+                        aliased_table = tables.get(str(chain_part))
+                        if aliased_table is not None:
+                            last_table = aliased_table
                         continue
 
                     # Ignore last chain part, it's likely an incomplete word or added characters

posthog/hogql/test/test_autocomplete.py

@@ -169,3 +169,33 @@ def test_autocomplete_incomplete_list(self):
         query = "select properties. from events"
         results = self._query_response(query=query, start=18, end=18)
         assert results.incomplete_list is True
+
+    def test_autocomplete_joined_tables(self):
+        query = "select p. from events e left join persons p on e.person_id = p.id"
+        results = self._query_response(query=query, start=9, end=9)
+
+        assert len(results.suggestions) != 0
+
+        keys = list(PERSONS_FIELDS.keys())
+
+        for index, key in enumerate(keys):
+            assert results.suggestions[index].label == key
+
+    def test_autocomplete_joined_table_contraints(self):
+        query = "select p.id from events e left join persons p on e.person_id = p."
+        results = self._query_response(query=query, start=65, end=65)
+
+        assert len(results.suggestions) != 0
+
+        keys = list(PERSONS_FIELDS.keys())
+
+        for index, key in enumerate(keys):
+            assert results.suggestions[index].label == key
+
+    def test_autocomplete_joined_tables_aliases(self):
+        query = "select  from events e left join persons p on e.person_id = p.id"
+        results = self._query_response(query=query, start=7, end=7)
+
+        assert len(results.suggestions) == 2
+        assert results.suggestions[0].label == "e"
+        assert results.suggestions[1].label == "p"