feat(data-warehouse): Override external data table fields

mypy-baseline.txt

@@ -97,7 +97,7 @@ posthog/hogql/database/database.py:0: error: "FieldOrTable" has no attribute "fi
 posthog/hogql/database/database.py:0: error: "FieldOrTable" has no attribute "fields"  [attr-defined]
 posthog/hogql/database/database.py:0: error: "FieldOrTable" has no attribute "fields"  [attr-defined]
 posthog/hogql/database/database.py:0: error: "FieldOrTable" has no attribute "fields"  [attr-defined]
-posthog/hogql/database/database.py:0: error: Incompatible types (expression has type "Literal['view', 'lazy_table']", TypedDict item "type" has type "Literal['integer', 'float', 'string', 'datetime', 'date', 'boolean', 'array', 'json', 'lazy_table', 'virtual_table', 'field_traverser']")  [typeddict-item]
+posthog/hogql/database/database.py:0: error: Incompatible types (expression has type "Literal['view', 'lazy_table']", TypedDict item "type" has type "Literal['integer', 'float', 'string', 'datetime', 'date', 'boolean', 'array', 'json', 'lazy_table', 'virtual_table', 'field_traverser', 'expression']")  [typeddict-item]
 posthog/warehouse/models/datawarehouse_saved_query.py:0: error: Argument 1 to "create_hogql_database" has incompatible type "int | None"; expected "int"  [arg-type]
 posthog/warehouse/models/datawarehouse_saved_query.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SelectQuery | SelectUnionQuery")  [assignment]
 posthog/models/user.py:0: error: Incompatible types in assignment (expression has type "None", base class "AbstractUser" defined the type as "CharField[str | int | Combinable, str]")  [assignment]
@@ -419,7 +419,6 @@ posthog/api/feature_flag.py:0: error: Item "Sequence[Any]" of "Any | Sequence[An
 posthog/api/feature_flag.py:0: error: Item "None" of "Any | Sequence[Any] | None" has no attribute "filters"  [union-attr]
 posthog/api/survey.py:0: error: Incompatible types in assignment (expression has type "Any | Sequence[Any] | None", variable has type "Survey | None")  [assignment]
 posthog/api/user.py:0: error: "User" has no attribute "social_auth"  [attr-defined]
-ee/api/role.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
 ee/api/dashboard_collaborator.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
 ee/api/test/base.py:0: error: Incompatible types in assignment (expression has type "None", variable has type "License")  [assignment]
 ee/api/test/base.py:0: error: "setUpTestData" undefined in superclass  [misc]
@@ -594,6 +593,8 @@ posthog/hogql/test/test_resolver.py:0: error: Incompatible types in assignment (
 posthog/hogql/test/test_resolver.py:0: error: "TestResolver" has no attribute "snapshot"  [attr-defined]
 posthog/hogql/test/test_resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SelectQuery")  [assignment]
 posthog/hogql/test/test_resolver.py:0: error: "TestResolver" has no attribute "snapshot"  [attr-defined]
+posthog/hogql/test/test_resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SelectQuery")  [assignment]
+posthog/hogql/test/test_resolver.py:0: error: "TestResolver" has no attribute "snapshot"  [attr-defined]
 posthog/hogql/test/test_resolver.py:0: error: Item "SelectUnionQueryType" of "SelectQueryType | SelectUnionQueryType | None" has no attribute "columns"  [union-attr]
 posthog/hogql/test/test_resolver.py:0: error: Item "None" of "SelectQueryType | SelectUnionQueryType | None" has no attribute "columns"  [union-attr]
 posthog/hogql/test/test_resolver.py:0: error: "FieldOrTable" has no attribute "fields"  [attr-defined]
@@ -649,7 +650,7 @@ posthog/hogql/functions/test/test_cohort.py:0: error: "TestCohort" has no attrib
 posthog/hogql/database/schema/test/test_channel_type.py:0: error: Value of type "list[Any] | None" is not indexable  [index]
 posthog/hogql/database/schema/test/test_channel_type.py:0: error: Value of type "list[Any] | None" is not indexable  [index]
 posthog/api/organization_member.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
-ee/api/feature_flag_role_access.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
+ee/api/role.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
 ee/clickhouse/views/insights.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
 posthog/queries/trends/test/test_person.py:0: error: "str" has no attribute "get"  [attr-defined]
 posthog/queries/trends/test/test_person.py:0: error: Invalid index type "int" for "HttpResponse"; expected type "str | bytes"  [index]
@@ -754,6 +755,7 @@ posthog/api/property_definition.py:0: error: Incompatible types in assignment (e
 posthog/api/property_definition.py:0: error: Item "AnonymousUser" of "User | AnonymousUser" has no attribute "organization"  [union-attr]
 posthog/api/property_definition.py:0: error: Item "None" of "Organization | Any | None" has no attribute "is_feature_available"  [union-attr]
 posthog/api/dashboards/dashboard_templates.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
+ee/api/feature_flag_role_access.py:0: error: Metaclass conflict: the metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its bases  [misc]
 posthog/temporal/tests/batch_exports/test_run_updates.py:0: error: Unused "type: ignore" comment  [unused-ignore]
 posthog/temporal/tests/batch_exports/test_run_updates.py:0: error: Unused "type: ignore" comment  [unused-ignore]
 posthog/temporal/tests/batch_exports/test_run_updates.py:0: error: Item "None" of "BatchExportRun | None" has no attribute "data_interval_start"  [union-attr]

posthog/hogql/autocomplete.py

@@ -1,7 +1,7 @@
 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
+from posthog.hogql.database.database import Database, create_hogql_database
 from posthog.hogql.database.models import (
     BooleanDatabaseField,
     DatabaseField,
@@ -239,6 +239,10 @@ def append_table_field_to_response(table: Table, suggestions: List[AutocompleteC
     details: List[str | None] = []
     table_fields = list(table.fields.items())
     for field_name, field_or_table in table_fields:
+        # Skip over hidden fields
+        if isinstance(field_or_table, ast.DatabaseField) and field_or_table.hidden:
+            continue
+
         keys.append(field_name)
         details.append(convert_field_or_table_to_type_string(field_or_table))
 
@@ -278,11 +282,17 @@ def extend_responses(
 
 
 # TODO: Support ast.SelectUnionQuery nodes
-def get_hogql_autocomplete(query: HogQLAutocomplete, team: Team) -> HogQLAutocompleteResponse:
+def get_hogql_autocomplete(
+    query: HogQLAutocomplete, team: Team, database_arg: Optional[Database] = None
+) -> HogQLAutocompleteResponse:
     response = HogQLAutocompleteResponse(suggestions=[], incomplete_list=False)
     timings = HogQLTimings()
 
-    database = create_hogql_database(team_id=team.pk, team_arg=team)
+    if database_arg is not None:
+        database = database_arg
+    else:
+        database = create_hogql_database(team_id=team.pk, team_arg=team)
+
     context = HogQLContext(team_id=team.pk, team=team, database=database)
 
     original_query_select = copy(query.select)

posthog/hogql/database/database.py

@@ -309,6 +309,7 @@ class _SerializedFieldBase(TypedDict):
         "lazy_table",
         "virtual_table",
         "field_traverser",
+        "expression",
     ]
 
 
@@ -346,6 +347,9 @@ def serialize_fields(field_input, context: HogQLContext) -> List[SerializedField
         if field_key == "team_id":
             pass
         elif isinstance(field, DatabaseField):
+            if field.hidden:
+                continue
+
             if isinstance(field, IntegerDatabaseField):
                 field_output.append({"key": field_key, "type": "integer"})
             elif isinstance(field, FloatDatabaseField):
@@ -362,6 +366,8 @@ def serialize_fields(field_input, context: HogQLContext) -> List[SerializedField
                 field_output.append({"key": field_key, "type": "json"})
             elif isinstance(field, StringArrayDatabaseField):
                 field_output.append({"key": field_key, "type": "array"})
+            elif isinstance(field, ExpressionField):
+                field_output.append({"key": field_key, "type": "expression"})
         elif isinstance(field, LazyJoin):
             is_view = isinstance(field.resolve_table(context), SavedQuery)
             field_output.append(

posthog/hogql/database/models.py

@@ -24,6 +24,7 @@ class DatabaseField(FieldOrTable):
     name: str
     array: Optional[bool] = None
     nullable: Optional[bool] = None
+    hidden: bool = False
 
 
 class IntegerDatabaseField(DatabaseField):
@@ -95,15 +96,11 @@ def get_asterisk(self):
         for key, field in self.fields.items():
             if key in fields_to_avoid:
                 continue
-            if (
-                isinstance(field, Table)
-                or isinstance(field, LazyJoin)
-                or isinstance(field, FieldTraverser)
-                or isinstance(field, ExpressionField)
-            ):
+            if isinstance(field, Table) or isinstance(field, LazyJoin) or isinstance(field, FieldTraverser):
                 pass  # ignore virtual tables and columns for now
             elif isinstance(field, DatabaseField):
-                asterisk[key] = field
+                if not field.hidden:  # Skip over hidden fields
+                    asterisk[key] = field
             else:
                 raise HogQLException(f"Unknown field type {type(field).__name__} for asterisk")
         return asterisk

posthog/hogql/database/test/__snapshots__/test_database.ambr

@@ -269,6 +269,14 @@
                   "distinct_id",
                   "person_id"
               ]
+          },
+          {
+              "key": "$virt_initial_referring_domain_type",
+              "type": "expression"
+          },
+          {
+              "key": "$virt_initial_channel_type",
+              "type": "expression"
           }
       ],
       "person_distinct_ids": [
@@ -1112,6 +1120,14 @@
                   "distinct_id",
                   "person_id"
               ]
+          },
+          {
+              "key": "$virt_initial_referring_domain_type",
+              "type": "expression"
+          },
+          {
+              "key": "$virt_initial_channel_type",
+              "type": "expression"
           }
       ],
       "person_distinct_ids": [

posthog/hogql/database/test/test_database.py

@@ -134,7 +134,7 @@ def test_database_expression_fields(self):
         query = print_ast(parse_select(sql), context, dialect="clickhouse")
         assert (
             query
-            == "SELECT number AS number FROM (SELECT numbers.number AS number FROM numbers(2) AS numbers) LIMIT 10000"
+            == "SELECT number AS number, expression AS expression, double AS double FROM (SELECT numbers.number AS number, plus(1, 1) AS expression, multiply(numbers.number, 2) AS double FROM numbers(2) AS numbers) LIMIT 10000"
         ), query
 
     def test_database_warehouse_joins(self):