fix: selecting from cohort people (#17610)

* fix: selecting from cohort people * add a test * include numbers table in python test * fix FE example for numbers * always include the person and cohort id * Update query snapshots * Update query snapshots --------- Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>
PostHog · Sep 26, 2023 · ec4faf2 · ec4faf2
1 parent 0b9e7ea
commit ec4faf2
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 5 deletions.
diff --git a/frontend/src/scenes/data-management/database/DatabaseTables.tsx b/frontend/src/scenes/data-management/database/DatabaseTables.tsx
@@ -88,7 +88,9 @@ export function DatabaseTables<T extends DatabaseSceneRow>({
                                               // TODO: Use `hogql` tag?
                                               query: `SELECT ${obj.columns
                                                   .filter(({ table, fields, chain }) => !table && !fields && !chain)
-                                                  .map(({ key }) => key)} FROM ${table} LIMIT 100`,
+                                                  .map(({ key }) => key)} FROM ${
+                                                  table === 'numbers' ? 'numbers(0, 10)' : table
+                                              } LIMIT 100`,
                                           },
                                       }
                                       return (

diff --git a/posthog/hogql/database/schema/cohort_people.py b/posthog/hogql/database/schema/cohort_people.py
@@ -23,11 +23,9 @@ def select_from_cohort_people_table(requested_fields: Dict[str, List[str]]):
 
     table_name = "raw_cohort_people"
 
+    # must always include the person and cohort ids regardless of what other fields are requested
     requested_fields = {"person_id": ["person_id"], "cohort_id": ["cohort_id"], **requested_fields}
-
-    fields: List[ast.Expr] = [
-        ast.Alias(alias=name, expr=ast.Field(chain=[table_name] + chain)) for name, chain in requested_fields.items()
-    ]
+    fields: List[ast.Expr] = [ast.Field(chain=[table_name] + chain) for name, chain in requested_fields.items()]
 
     return ast.SelectQuery(
         select=fields,

diff --git a/posthog/hogql/database/test/test_database.py b/posthog/hogql/database/test/test_database.py
@@ -4,6 +4,7 @@
 from unittest.mock import patch
 import pytest
 from django.test import override_settings
+from parameterized import parameterized
 
 from posthog.hogql.database.database import create_hogql_database, serialize_database
 from posthog.test.base import BaseTest
@@ -26,6 +27,21 @@ def test_serialize_database_with_person_on_events_enabled(self):
             serialized_database = serialize_database(create_hogql_database(team_id=self.team.pk))
             assert json.dumps(serialized_database, indent=4) == self.snapshot
 
+    @parameterized.expand([False, True])
+    def test_can_select_from_each_table_at_all(self, poe_enabled: bool) -> None:
+        with override_settings(PERSON_ON_EVENTS_OVERRIDE=poe_enabled):
+            serialized_database = serialize_database(create_hogql_database(team_id=self.team.pk))
+            for table, possible_columns in serialized_database.items():
+                if table == "numbers":
+                    execute_hogql_query("SELECT number FROM numbers(10) LIMIT 100", self.team)
+                else:
+                    columns = [
+                        x["key"]
+                        for x in possible_columns
+                        if "table" not in x and "chain" not in x and "fields" not in x
+                    ]
+                    execute_hogql_query(f"SELECT {','.join(columns)} FROM {table}", team=self.team)
+
     @patch("posthog.hogql.query.sync_execute", return_value=(None, None))
     @pytest.mark.usefixtures("unittest_snapshot")
     def test_database_with_warehouse_tables(self, patch_execute):