feat(autocomplete): Support field traverser nodes (#20287)

* Support field traverser nodes

* Fixed mypy

* Added the incomplete field to the response

* Remove test

frontend/src/queries/nodes/HogQLQuery/HogQLQueryEditor.tsx

@@ -265,6 +265,7 @@ export function HogQLQueryEditor(props: HogQLQueryEditorProps): JSX.Element {
 
                                         return {
                                             suggestions,
+                                            incomplete: response.incomplete_list,
                                         }
                                     },
                                 })

frontend/src/queries/schema.json

@@ -1938,14 +1938,18 @@
         "HogQLAutocompleteResponse": {
             "additionalProperties": false,
             "properties": {
+                "incomplete_list": {
+                    "description": "Whether or not the suggestions returned are complete",
+                    "type": "boolean"
+                },
                 "suggestions": {
                     "items": {
                         "$ref": "#/definitions/AutocompleteCompletionItem"
                     },
                     "type": "array"
                 }
             },
-            "required": ["suggestions"],
+            "required": ["suggestions", "incomplete_list"],
             "type": "object"
         },
         "HogQLExpression": {
@@ -3572,14 +3576,18 @@
                 {
                     "additionalProperties": false,
                     "properties": {
+                        "incomplete_list": {
+                            "description": "Whether or not the suggestions returned are complete",
+                            "type": "boolean"
+                        },
                         "suggestions": {
                             "items": {
                                 "$ref": "#/definitions/AutocompleteCompletionItem"
                             },
                             "type": "array"
                         }
                     },
-                    "required": ["suggestions"],
+                    "required": ["suggestions", "incomplete_list"],
                     "type": "object"
                 },
                 {

frontend/src/queries/schema.ts

@@ -296,6 +296,8 @@ export interface AutocompleteCompletionItem {
 
 export interface HogQLAutocompleteResponse {
     suggestions: AutocompleteCompletionItem[]
+    /** Whether or not the suggestions returned are complete */
+    incomplete_list: boolean
 }
 
 export interface HogQLMetadata extends DataNode {

mypy-baseline.txt

@@ -364,6 +364,7 @@ posthog/hogql/query.py:0: error: Incompatible types in assignment (expression ha
 posthog/hogql/query.py:0: error: Argument 1 to "get_default_limit_for_context" has incompatible type "LimitContext | None"; expected "LimitContext"  [arg-type]
 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_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/autocomplete.py

@@ -1,4 +1,4 @@
-from copy import copy
+from copy import copy, deepcopy
 from typing import Callable, Dict, List, Optional, cast
 from posthog.hogql.context import HogQLContext
 from posthog.hogql.database.database import create_hogql_database
@@ -82,7 +82,7 @@ def constant_type_to_database_field(constant_type: ConstantType, name: str) -> D
     return DatabaseField(name=name)
 
 
-def convert_field_or_table_to_type_string(field_or_table: FieldOrTable) -> str:
+def convert_field_or_table_to_type_string(field_or_table: FieldOrTable) -> str | None:
     if isinstance(field_or_table, ast.BooleanDatabaseField):
         return "Boolean"
     if isinstance(field_or_table, ast.IntegerDatabaseField):
@@ -100,7 +100,7 @@ def convert_field_or_table_to_type_string(field_or_table: FieldOrTable) -> str:
     if isinstance(field_or_table, (ast.Table, ast.LazyJoin)):
         return "Table"
 
-    return ""
+    return None
 
 
 def get_table(context: HogQLContext, join_expr: ast.JoinExpr, ctes: Optional[Dict[str, CTE]]) -> None | Table:
@@ -181,6 +181,55 @@ def to_printed_hogql(self):
     return None
 
 
+# Replaces all ast.FieldTraverser with the underlying node
+def resolve_table_field_traversers(table: Table) -> Table:
+    new_table = deepcopy(table)
+    new_fields: Dict[str, FieldOrTable] = {}
+    for key, field in list(new_table.fields.items()):
+        if not isinstance(field, ast.FieldTraverser):
+            new_fields[key] = field
+            continue
+
+        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)
+            elif isinstance(current_table_or_field, LazyJoin):
+                chain_field = current_table_or_field.join_table.fields.get(chain)
+            elif isinstance(current_table_or_field, DatabaseField):
+                chain_field = current_table_or_field
+            else:
+                # Cant find the field, default back
+                new_fields[key] = field
+                break
+
+            if chain_field is not None:
+                current_table_or_field = chain_field
+                new_fields[key] = chain_field
+
+    new_table.fields = new_fields
+    return new_table
+
+
+def append_table_field_to_response(table: Table, suggestions: List[AutocompleteCompletionItem]) -> None:
+    keys: List[str] = []
+    details: List[str | None] = []
+    table_fields = list(table.fields.items())
+    for field_name, field_or_table in table_fields:
+        keys.append(field_name)
+        details.append(convert_field_or_table_to_type_string(field_or_table))
+
+    extend_responses(keys=keys, suggestions=suggestions, details=details)
+
+    available_functions = ALL_EXPOSED_FUNCTION_NAMES
+    extend_responses(
+        available_functions,
+        suggestions,
+        Kind.Function,
+        insert_text=lambda key: f"{key}()",
+    )
+
+
 def extend_responses(
     keys: List[str],
     suggestions: List[AutocompleteCompletionItem],
@@ -207,7 +256,7 @@ def extend_responses(
 
 # TODO: Support ast.SelectUnionQuery nodes
 def get_hogql_autocomplete(query: HogQLAutocomplete, team: Team) -> HogQLAutocompleteResponse:
-    response = HogQLAutocompleteResponse(suggestions=[])
+    response = HogQLAutocompleteResponse(suggestions=[], incomplete_list=False)
 
     database = create_hogql_database(team_id=team.pk, team_arg=team)
     context = HogQLContext(team_id=team.pk, team=team, database=database)
@@ -259,8 +308,6 @@ def get_hogql_autocomplete(query: HogQLAutocomplete, team: Team) -> HogQLAutocom
                 and nearest_select.select_from is not None
                 and not isinstance(parent_node, ast.JoinExpr)
             ):
-                # TODO: add logic for FieldTraverser field types
-
                 # Handle fields
                 table = get_table(context, nearest_select.select_from, ctes)
                 if table is None:
@@ -269,9 +316,17 @@ 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:
-                        extend_responses([str(chain_part)], response.suggestions, Kind.Folder)
+                        alias = nearest_select.select_from.alias
+                        assert alias is not None
+                        extend_responses(
+                            keys=[alias],
+                            suggestions=response.suggestions,
+                            kind=Kind.Folder,
+                            details=["Table"],
+                        )
                         break
 
                     if table_has_alias and index == 0:
@@ -280,24 +335,12 @@ def get_hogql_autocomplete(query: HogQLAutocomplete, team: Team) -> HogQLAutocom
                     # Ignore last chain part, it's likely an incomplete word or added characters
                     is_last_part = index >= (chain_len - 2)
 
+                    # Replaces all ast.FieldTraverser with the underlying node
+                    last_table = resolve_table_field_traversers(last_table)
+
                     if is_last_part:
                         if last_table.fields.get(str(chain_part)) is None:
-                            keys: List[str] = []
-                            details: List[str | None] = []
-                            table_fields = list(table.fields.items())
-                            for field_name, field_or_table in table_fields:
-                                keys.append(field_name)
-                                details.append(convert_field_or_table_to_type_string(field_or_table))
-
-                            extend_responses(keys=keys, suggestions=response.suggestions, details=details)
-
-                            available_functions = ALL_EXPOSED_FUNCTION_NAMES
-                            extend_responses(
-                                available_functions,
-                                response.suggestions,
-                                Kind.Function,
-                                insert_text=lambda key: f"{key}()",
-                            )
+                            append_table_field_to_response(table=last_table, suggestions=response.suggestions)
                             break
 
                         field = last_table.fields[str(chain_part)]
@@ -328,12 +371,22 @@ def get_hogql_autocomplete(query: HogQLAutocomplete, team: Team) -> HogQLAutocom
                                     suggestions=response.suggestions,
                                     details=[prop["property_type"] for prop in properties],
                                 )
+                                response.incomplete_list = True
                         elif isinstance(field, VirtualTable) or isinstance(field, LazyTable):
-                            fields = list(last_table.fields.keys())
-                            extend_responses(fields, response.suggestions)
+                            fields = list(last_table.fields.items())
+                            extend_responses(
+                                keys=[key for key, field in fields],
+                                suggestions=response.suggestions,
+                                details=[convert_field_or_table_to_type_string(field) for key, field in fields],
+                            )
                         elif isinstance(field, LazyJoin):
-                            fields = list(field.join_table.fields.keys())
-                            extend_responses(fields, response.suggestions)
+                            fields = list(field.join_table.fields.items())
+
+                            extend_responses(
+                                keys=[key for key, field in fields],
+                                suggestions=response.suggestions,
+                                details=[convert_field_or_table_to_type_string(field) for key, field in fields],
+                            )
                         break
                     else:
                         field = last_table.fields[str(chain_part)]
@@ -345,7 +398,12 @@ def get_hogql_autocomplete(query: HogQLAutocomplete, team: Team) -> HogQLAutocom
                 # Handle table names
                 if len(node.chain) == 1:
                     table_names = database.get_all_tables()
-                    extend_responses(table_names, response.suggestions, Kind.Folder)
+                    extend_responses(
+                        keys=table_names,
+                        suggestions=response.suggestions,
+                        kind=Kind.Folder,
+                        details=["Table"] * len(table_names),  # type: ignore
+                    )
         except Exception:
             pass
 

posthog/hogql/test/test_autocomplete.py

@@ -148,3 +148,24 @@ def test_autocomplete_cte_constant_type(self):
         assert results.suggestions[0].label == "potato"
         assert "event" not in [suggestion.label for suggestion in results.suggestions]
         assert "properties" not in [suggestion.label for suggestion in results.suggestions]
+
+    def test_autocomplete_field_traversers(self):
+        query = "select person. from events"
+        results = self._query_response(query=query, start=14, end=14)
+        assert len(results.suggestions) != 0
+
+    def test_autocomplete_table_alias(self):
+        query = "select  from events e"
+        results = self._query_response(query=query, start=7, end=7)
+        assert len(results.suggestions) != 0
+        assert results.suggestions[0].label == "e"
+
+    def test_autocomplete_complete_list(self):
+        query = "select event from events"
+        results = self._query_response(query=query, start=7, end=12)
+        assert results.incomplete_list is False
+
+    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

posthog/schema.py

@@ -328,6 +328,7 @@ class HogQLAutocompleteResponse(BaseModel):
     model_config = ConfigDict(
         extra="forbid",
     )
+    incomplete_list: bool = Field(..., description="Whether or not the suggestions returned are complete")
     suggestions: List[AutocompleteCompletionItem]
 
 
@@ -667,6 +668,7 @@ class QueryResponseAlternative9(BaseModel):
     model_config = ConfigDict(
         extra="forbid",
     )
+    incomplete_list: bool = Field(..., description="Whether or not the suggestions returned are complete")
     suggestions: List[AutocompleteCompletionItem]