From 42bf36683078150b871cc71487765841c33241be Mon Sep 17 00:00:00 2001 From: Michael Matloka Date: Tue, 19 Mar 2024 19:40:30 +0100 Subject: [PATCH 1/8] fix(insights): Attribute async queries to users --- posthog/api/query.py | 1 + posthog/clickhouse/client/execute_async.py | 39 ++++---- .../client/test/test_execute_async.py | 62 ++++++++---- posthog/schema.py | 96 +++++++------------ posthog/tasks/tasks.py | 8 +- 5 files changed, 109 insertions(+), 97 deletions(-) diff --git a/posthog/api/query.py b/posthog/api/query.py index d8f45531253a0..19b70677397e3 100644 --- a/posthog/api/query.py +++ b/posthog/api/query.py @@ -67,6 +67,7 @@ def create(self, request, *args, **kwargs) -> Response: if data.async_: query_status = enqueue_process_query_task( team_id=self.team.pk, + user_id=self.request.user.pk, query_json=request.data["query"], query_id=client_query_id, refresh_requested=data.refresh, diff --git a/posthog/clickhouse/client/execute_async.py b/posthog/clickhouse/client/execute_async.py index 06f7fc639f824..fc7e2cc4c84c6 100644 --- a/posthog/clickhouse/client/execute_async.py +++ b/posthog/clickhouse/client/execute_async.py @@ -1,5 +1,6 @@ import datetime import json +from typing import Optional import uuid import structlog @@ -69,11 +70,7 @@ def delete_query_status(self): def execute_process_query( - team_id, - query_id, - query_json, - limit_context, - refresh_requested, + team_id: int, user_id: int, query_id: str, query_json: str, limit_context: LimitContext, refresh_requested: bool ): manager = QueryStatusManager(query_id, team_id) @@ -91,7 +88,7 @@ def execute_process_query( QUERY_WAIT_TIME.observe(wait_duration) try: - tag_queries(client_query_id=query_id, team_id=team_id) + tag_queries(client_query_id=query_id, team_id=team_id, user_id=user_id) results = process_query( team=team, query_json=query_json, limit_context=limit_context, refresh_requested=refresh_requested ) @@ -113,12 +110,13 @@ def execute_process_query( def enqueue_process_query_task( - team_id, - query_json, - query_id=None, - refresh_requested=False, - bypass_celery=False, - force=False, + team_id: int, + user_id: int, + query_json: str, + query_id: Optional[str] = None, + refresh_requested: bool = False, + force: bool = False, + _test_only_bypass_celery: bool = False, ) -> QueryStatus: if not query_id: query_id = uuid.uuid4().hex @@ -136,14 +134,23 @@ def enqueue_process_query_task( query_status = QueryStatus(id=query_id, team_id=team_id, start_time=datetime.datetime.now(datetime.timezone.utc)) manager.store_query_status(query_status) - if bypass_celery: - # Call directly ( for testing ) + if _test_only_bypass_celery: process_query_task( - team_id, query_id, query_json, limit_context=LimitContext.QUERY_ASYNC, refresh_requested=refresh_requested + team_id, + user_id, + query_id, + query_json, + limit_context=LimitContext.QUERY_ASYNC, + refresh_requested=refresh_requested, ) else: task = process_query_task.delay( - team_id, query_id, query_json, limit_context=LimitContext.QUERY_ASYNC, refresh_requested=refresh_requested + team_id, + user_id, + query_id, + query_json, + limit_context=LimitContext.QUERY_ASYNC, + refresh_requested=refresh_requested, ) query_status.task_id = task.id manager.store_query_status(query_status) diff --git a/posthog/clickhouse/client/test/test_execute_async.py b/posthog/clickhouse/client/test/test_execute_async.py index 0d7a7281e6a4b..085e7708b9232 100644 --- a/posthog/clickhouse/client/test/test_execute_async.py +++ b/posthog/clickhouse/client/test/test_execute_async.py @@ -24,6 +24,7 @@ def setUp(self): self.organization = Organization.objects.create(name="test") self.team = Team.objects.create(organization=self.organization) self.team_id = self.team.pk + self.user_id = 1337 self.query_id = "test_query_id" self.query_json = {} self.limit_context = None @@ -41,7 +42,9 @@ def test_execute_process_query(self, mock_process_query, mock_redis_client): mock_process_query.return_value = [float("inf"), float("-inf"), float("nan"), 1.0, "👍"] - execute_process_query(self.team_id, self.query_id, self.query_json, self.limit_context, self.refresh_requested) + execute_process_query( + self.team_id, self.user_id, self.query_id, self.query_json, self.limit_context, self.refresh_requested + ) mock_redis_client.assert_called_once() mock_process_query.assert_called_once() @@ -55,15 +58,16 @@ def test_execute_process_query(self, mock_process_query, mock_redis_client): class ClickhouseClientTestCase(TestCase, ClickhouseTestMixin): def setUp(self): - self.organization = Organization.objects.create(name="test") - self.team = Team.objects.create(organization=self.organization) - self.team_id = self.team.pk + self.organization: Organization = Organization.objects.create(name="test") + self.team: Team = Team.objects.create(organization=self.organization) + self.team_id: int = self.team.pk + self.user_id: int = 2137 @snapshot_clickhouse_queries def test_async_query_client(self): query = build_query("SELECT 1+1") team_id = self.team_id - query_id = client.enqueue_process_query_task(team_id, query, bypass_celery=True).id + query_id = client.enqueue_process_query_task(team_id, self.user_id, query, _test_only_bypass_celery=True).id result = client.get_query_status(team_id, query_id) self.assertFalse(result.error, result.error_message) self.assertTrue(result.complete) @@ -74,11 +78,13 @@ def test_async_query_client_errors(self): self.assertRaises( HogQLException, client.enqueue_process_query_task, - **{"team_id": (self.team_id), "query_json": query, "bypass_celery": True}, + **{"team_id": self.team_id, "user_id": self.user_id, "query_json": query, "_test_only_bypass_celery": True}, ) query_id = uuid.uuid4().hex try: - client.enqueue_process_query_task(self.team_id, query, query_id=query_id, bypass_celery=True) + client.enqueue_process_query_task( + self.team_id, self.user_id, query, query_id=query_id, _test_only_bypass_celery=True + ) except Exception: pass @@ -89,7 +95,7 @@ def test_async_query_client_errors(self): def test_async_query_client_uuid(self): query = build_query("SELECT toUUID('00000000-0000-0000-0000-000000000000')") team_id = self.team_id - query_id = client.enqueue_process_query_task(team_id, query, bypass_celery=True).id + query_id = client.enqueue_process_query_task(team_id, self.user_id, query, _test_only_bypass_celery=True).id result = client.get_query_status(team_id, query_id) self.assertFalse(result.error, result.error_message) self.assertTrue(result.complete) @@ -99,7 +105,7 @@ def test_async_query_client_does_not_leak(self): query = build_query("SELECT 1+1") team_id = self.team_id wrong_team = 5 - query_id = client.enqueue_process_query_task(team_id, query, bypass_celery=True).id + query_id = client.enqueue_process_query_task(team_id, self.user_id, query, _test_only_bypass_celery=True).id try: client.get_query_status(wrong_team, query_id) @@ -111,13 +117,19 @@ def test_async_query_client_is_lazy(self, execute_sync_mock): query = build_query("SELECT 4 + 4") query_id = uuid.uuid4().hex team_id = self.team_id - client.enqueue_process_query_task(team_id, query, query_id=query_id, bypass_celery=True) + client.enqueue_process_query_task( + team_id, self.user_id, query, query_id=query_id, _test_only_bypass_celery=True + ) # Try the same query again - client.enqueue_process_query_task(team_id, query, query_id=query_id, bypass_celery=True) + client.enqueue_process_query_task( + team_id, self.user_id, query, query_id=query_id, _test_only_bypass_celery=True + ) # Try the same query again (for good measure!) - client.enqueue_process_query_task(team_id, query, query_id=query_id, bypass_celery=True) + client.enqueue_process_query_task( + team_id, self.user_id, query, query_id=query_id, _test_only_bypass_celery=True + ) # Assert that we only called clickhouse once execute_sync_mock.assert_called_once() @@ -127,13 +139,19 @@ def test_async_query_client_is_lazy_but_not_too_lazy(self, execute_sync_mock): query = build_query("SELECT 8 + 8") query_id = uuid.uuid4().hex team_id = self.team_id - client.enqueue_process_query_task(team_id, query, query_id=query_id, bypass_celery=True) + client.enqueue_process_query_task( + team_id, self.user_id, query, query_id=query_id, _test_only_bypass_celery=True + ) # Try the same query again, but with force - client.enqueue_process_query_task(team_id, query, query_id=query_id, bypass_celery=True, force=True) + client.enqueue_process_query_task( + team_id, self.user_id, query, query_id=query_id, _test_only_bypass_celery=True, force=True + ) # Try the same query again (for good measure!) - client.enqueue_process_query_task(team_id, query, query_id=query_id, bypass_celery=True) + client.enqueue_process_query_task( + team_id, self.user_id, query, query_id=query_id, _test_only_bypass_celery=True + ) # Assert that we called clickhouse twice self.assertEqual(execute_sync_mock.call_count, 2) @@ -145,13 +163,19 @@ def test_async_query_client_manual_query_uuid(self, execute_sync_mock): query = build_query("SELECT 8 + 8") team_id = self.team_id query_id = "I'm so unique" - client.enqueue_process_query_task(team_id, query, query_id=query_id, bypass_celery=True) + client.enqueue_process_query_task( + team_id, self.user_id, query, query_id=query_id, _test_only_bypass_celery=True + ) # Try the same query again, but with force - client.enqueue_process_query_task(team_id, query, query_id=query_id, bypass_celery=True, force=True) + client.enqueue_process_query_task( + team_id, self.user_id, query, query_id=query_id, _test_only_bypass_celery=True, force=True + ) # Try the same query again (for good measure!) - client.enqueue_process_query_task(team_id, query, query_id=query_id, bypass_celery=True) + client.enqueue_process_query_task( + team_id, self.user_id, query, query_id=query_id, _test_only_bypass_celery=True + ) # Assert that we called clickhouse twice self.assertEqual(execute_sync_mock.call_count, 2) @@ -186,4 +210,4 @@ def test_client_strips_comments_from_request(self): # Make sure it still includes the "annotation" comment that includes # request routing information for debugging purposes - self.assertIn("/* request:1 */", first_query) + self.assertIn(f"/* user_id:{self.user_id} request:1 */", first_query) diff --git a/posthog/schema.py b/posthog/schema.py index dc77da163db17..5326b78126e65 100644 --- a/posthog/schema.py +++ b/posthog/schema.py @@ -4,7 +4,7 @@ from __future__ import annotations from enum import Enum -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Any, Dict, List, Literal, Optional, TypeAlias, Union from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, RootModel @@ -3001,6 +3001,38 @@ class HogQLMetadata(BaseModel): table: Optional[str] = Field(default=None, description="Table to validate the expression against") +AnyQuery: TypeAlias = Union[ + EventsNode, + ActionsNode, + PersonsNode, + DataWarehouseNode, + TimeToSeeDataSessionsQuery, + EventsQuery, + ActorsQuery, + InsightActorsQuery, + InsightActorsQueryOptions, + SessionsTimelineQuery, + HogQLQuery, + HogQLMetadata, + HogQLAutocomplete, + WebOverviewQuery, + WebStatsTableQuery, + WebTopClicksQuery, + DataVisualizationNode, + DataTableNode, + SavedInsightNode, + InsightVizNode, + TrendsQuery, + FunnelsQuery, + RetentionQuery, + PathsQuery, + StickinessQuery, + LifecycleQuery, + FunnelCorrelationQuery, + DatabaseSchemaQuery, +] + + class QueryRequest(BaseModel): model_config = ConfigDict( extra="forbid", @@ -3014,36 +3046,7 @@ class QueryRequest(BaseModel): client_query_id: Optional[str] = Field( default=None, description="Client provided query ID. Can be used to retrieve the status or cancel the query." ) - query: Union[ - EventsNode, - ActionsNode, - PersonsNode, - DataWarehouseNode, - TimeToSeeDataSessionsQuery, - EventsQuery, - ActorsQuery, - InsightActorsQuery, - InsightActorsQueryOptions, - SessionsTimelineQuery, - HogQLQuery, - HogQLMetadata, - HogQLAutocomplete, - WebOverviewQuery, - WebStatsTableQuery, - WebTopClicksQuery, - DataVisualizationNode, - DataTableNode, - SavedInsightNode, - InsightVizNode, - TrendsQuery, - FunnelsQuery, - RetentionQuery, - PathsQuery, - StickinessQuery, - LifecycleQuery, - FunnelCorrelationQuery, - DatabaseSchemaQuery, - ] = Field( + query: AnyQuery = Field( ..., description='Submit a JSON string representing a query for PostHog data analysis, for example a HogQL query.\n\nExample payload:\n\n```\n\n{"query": {"kind": "HogQLQuery", "query": "select * from events limit 100"}}\n\n```\n\nFor more details on HogQL queries, see the [PostHog HogQL documentation](/docs/hogql#api-access).', discriminator="kind", @@ -3085,36 +3088,7 @@ class QuerySchemaRoot( ] ] ): - root: Union[ - EventsNode, - ActionsNode, - PersonsNode, - DataWarehouseNode, - TimeToSeeDataSessionsQuery, - EventsQuery, - ActorsQuery, - InsightActorsQuery, - InsightActorsQueryOptions, - SessionsTimelineQuery, - HogQLQuery, - HogQLMetadata, - HogQLAutocomplete, - WebOverviewQuery, - WebStatsTableQuery, - WebTopClicksQuery, - DataVisualizationNode, - DataTableNode, - SavedInsightNode, - InsightVizNode, - TrendsQuery, - FunnelsQuery, - RetentionQuery, - PathsQuery, - StickinessQuery, - LifecycleQuery, - FunnelCorrelationQuery, - DatabaseSchemaQuery, - ] = Field(..., discriminator="kind") + root: AnyQuery = Field(..., discriminator="kind") PropertyGroupFilterValue.model_rebuild() diff --git a/posthog/tasks/tasks.py b/posthog/tasks/tasks.py index 5eff6afd33fe2..721eb92ee31ac 100644 --- a/posthog/tasks/tasks.py +++ b/posthog/tasks/tasks.py @@ -33,7 +33,12 @@ def redis_heartbeat() -> None: @shared_task(ignore_result=True, queue=CeleryQueue.ANALYTICS_QUERIES.value) def process_query_task( - team_id: str, query_id: str, query_json: Any, limit_context: Any = None, refresh_requested: bool = False + team_id: str, + user_id: str, + query_id: str, + query_json: Any, + limit_context: Any = None, + refresh_requested: bool = False, ) -> None: """ Kick off query @@ -43,6 +48,7 @@ def process_query_task( execute_process_query( team_id=team_id, + user_id=user_id, query_id=query_id, query_json=query_json, limit_context=limit_context, From 75b59e18daea3d3ce6dabc003551cf2db91207c3 Mon Sep 17 00:00:00 2001 From: Michael Matloka Date: Tue, 19 Mar 2024 20:00:50 +0100 Subject: [PATCH 2/8] Fix re-raising ClickHouse errors from threads --- posthog/errors.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/posthog/errors.py b/posthog/errors.py index 39f07be762d00..6b044497bc1ad 100644 --- a/posthog/errors.py +++ b/posthog/errors.py @@ -1,6 +1,6 @@ from dataclasses import dataclass import re -from typing import Dict +from typing import Dict, Optional from clickhouse_driver.errors import ServerException @@ -8,9 +8,10 @@ class InternalCHQueryError(ServerException): - code_name: str + code_name: Optional[str] + """Can be null if re-raised from a thread (see `failhard_threadhook_context`).""" - def __init__(self, message, *, code=None, nested=None, code_name): + def __init__(self, message, *, code=None, nested=None, code_name=None): self.code_name = code_name super().__init__(message, code, nested) From 92bfec6615a5869f093edfd04b37ce055c275e61 Mon Sep 17 00:00:00 2001 From: Michael Matloka Date: Tue, 19 Mar 2024 20:11:17 +0100 Subject: [PATCH 3/8] Improve typing --- posthog/api/query.py | 2 +- posthog/clickhouse/client/execute_async.py | 9 +++++++-- posthog/tasks/tasks.py | 9 +++++---- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/posthog/api/query.py b/posthog/api/query.py index 19b70677397e3..e30853655c749 100644 --- a/posthog/api/query.py +++ b/posthog/api/query.py @@ -70,7 +70,7 @@ def create(self, request, *args, **kwargs) -> Response: user_id=self.request.user.pk, query_json=request.data["query"], query_id=client_query_id, - refresh_requested=data.refresh, + refresh_requested=data.refresh or False, ) return Response(query_status.model_dump(), status=status.HTTP_202_ACCEPTED) diff --git a/posthog/clickhouse/client/execute_async.py b/posthog/clickhouse/client/execute_async.py index fc7e2cc4c84c6..4671b0060299b 100644 --- a/posthog/clickhouse/client/execute_async.py +++ b/posthog/clickhouse/client/execute_async.py @@ -70,7 +70,12 @@ def delete_query_status(self): def execute_process_query( - team_id: int, user_id: int, query_id: str, query_json: str, limit_context: LimitContext, refresh_requested: bool + team_id: int, + user_id: int, + query_id: str, + query_json: dict, + limit_context: Optional[LimitContext], + refresh_requested: bool, ): manager = QueryStatusManager(query_id, team_id) @@ -112,7 +117,7 @@ def execute_process_query( def enqueue_process_query_task( team_id: int, user_id: int, - query_json: str, + query_json: dict, query_id: Optional[str] = None, refresh_requested: bool = False, force: bool = False, diff --git a/posthog/tasks/tasks.py b/posthog/tasks/tasks.py index 721eb92ee31ac..78895c28d7b90 100644 --- a/posthog/tasks/tasks.py +++ b/posthog/tasks/tasks.py @@ -1,5 +1,5 @@ import time -from typing import Any, Optional +from typing import Optional from uuid import UUID from celery import shared_task @@ -9,6 +9,7 @@ from prometheus_client import Gauge from posthog.cloud_utils import is_cloud +from posthog.hogql.constants import LimitContext from posthog.metrics import pushed_metrics_registry from posthog.ph_client import get_ph_client from posthog.redis import get_client @@ -33,11 +34,11 @@ def redis_heartbeat() -> None: @shared_task(ignore_result=True, queue=CeleryQueue.ANALYTICS_QUERIES.value) def process_query_task( - team_id: str, + team_id: int, user_id: str, query_id: str, - query_json: Any, - limit_context: Any = None, + query_json: dict, + limit_context: Optional[LimitContext] = None, refresh_requested: bool = False, ) -> None: """ From b96a5ba1765fc575ee883d12c5d1c8c16dede881 Mon Sep 17 00:00:00 2001 From: github-actions <41898282+github-actions[bot]@users.noreply.github.com> Date: Tue, 19 Mar 2024 19:25:48 +0000 Subject: [PATCH 4/8] Update query snapshots --- .../test/__snapshots__/test_session_recordings.ambr | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/posthog/session_recordings/test/__snapshots__/test_session_recordings.ambr b/posthog/session_recordings/test/__snapshots__/test_session_recordings.ambr index fad3c08168d0b..c2982e4f8526f 100644 --- a/posthog/session_recordings/test/__snapshots__/test_session_recordings.ambr +++ b/posthog/session_recordings/test/__snapshots__/test_session_recordings.ambr @@ -778,7 +778,7 @@ FROM "posthog_persondistinctid" INNER JOIN "posthog_person" ON ("posthog_persondistinctid"."person_id" = "posthog_person"."id") WHERE ("posthog_persondistinctid"."distinct_id" IN ('user2', - 'user_one_2') + 'user_one_0') AND "posthog_persondistinctid"."team_id" = 2) /*controller='project_session_recordings-list',route='api/projects/%28%3FP%3Cparent_lookup_team_id%3E%5B%5E/.%5D%2B%29/session_recordings/%3F%24'*/ ''' # --- From da6cd2876282d8e098bbbd85b900a088dc7b7330 Mon Sep 17 00:00:00 2001 From: Michael Matloka Date: Tue, 19 Mar 2024 20:43:25 +0100 Subject: [PATCH 5/8] Fix more typing --- posthog/tasks/tasks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/posthog/tasks/tasks.py b/posthog/tasks/tasks.py index 78895c28d7b90..bead27cbd1eec 100644 --- a/posthog/tasks/tasks.py +++ b/posthog/tasks/tasks.py @@ -35,7 +35,7 @@ def redis_heartbeat() -> None: @shared_task(ignore_result=True, queue=CeleryQueue.ANALYTICS_QUERIES.value) def process_query_task( team_id: int, - user_id: str, + user_id: int, query_id: str, query_json: dict, limit_context: Optional[LimitContext] = None, From 15df6ffd00d906b6f45f73d9b3f53ec066d0c9d1 Mon Sep 17 00:00:00 2001 From: github-actions <41898282+github-actions[bot]@users.noreply.github.com> Date: Tue, 19 Mar 2024 19:51:28 +0000 Subject: [PATCH 6/8] Update query snapshots --- .../test/__snapshots__/test_session_recordings.ambr | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/posthog/session_recordings/test/__snapshots__/test_session_recordings.ambr b/posthog/session_recordings/test/__snapshots__/test_session_recordings.ambr index d5a3675a2c2ca..10b7a257b2a43 100644 --- a/posthog/session_recordings/test/__snapshots__/test_session_recordings.ambr +++ b/posthog/session_recordings/test/__snapshots__/test_session_recordings.ambr @@ -778,7 +778,7 @@ FROM "posthog_persondistinctid" INNER JOIN "posthog_person" ON ("posthog_persondistinctid"."person_id" = "posthog_person"."id") WHERE ("posthog_persondistinctid"."distinct_id" IN ('user2', - 'user_one_0') + 'user_one_2') AND "posthog_persondistinctid"."team_id" = 2) /*controller='project_session_recordings-list',route='api/projects/%28%3FP%3Cparent_lookup_team_id%3E%5B%5E/.%5D%2B%29/session_recordings/%3F%24'*/ ''' # --- From 6f24fb60b53d7f1fffe7bf5cf248548dd6dcb5d6 Mon Sep 17 00:00:00 2001 From: github-actions <41898282+github-actions[bot]@users.noreply.github.com> Date: Tue, 19 Mar 2024 20:13:13 +0000 Subject: [PATCH 7/8] Update query snapshots --- .../test/__snapshots__/test_session_recordings.ambr | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/posthog/session_recordings/test/__snapshots__/test_session_recordings.ambr b/posthog/session_recordings/test/__snapshots__/test_session_recordings.ambr index 10b7a257b2a43..d5a3675a2c2ca 100644 --- a/posthog/session_recordings/test/__snapshots__/test_session_recordings.ambr +++ b/posthog/session_recordings/test/__snapshots__/test_session_recordings.ambr @@ -778,7 +778,7 @@ FROM "posthog_persondistinctid" INNER JOIN "posthog_person" ON ("posthog_persondistinctid"."person_id" = "posthog_person"."id") WHERE ("posthog_persondistinctid"."distinct_id" IN ('user2', - 'user_one_2') + 'user_one_0') AND "posthog_persondistinctid"."team_id" = 2) /*controller='project_session_recordings-list',route='api/projects/%28%3FP%3Cparent_lookup_team_id%3E%5B%5E/.%5D%2B%29/session_recordings/%3F%24'*/ ''' # --- From d5166488c96f2cd5900d45ca968106fecc888d19 Mon Sep 17 00:00:00 2001 From: Michael Matloka Date: Tue, 19 Mar 2024 23:57:08 +0100 Subject: [PATCH 8/8] Update schema.py --- posthog/schema.py | 96 ++++++++++++++++++++++++++++++----------------- 1 file changed, 61 insertions(+), 35 deletions(-) diff --git a/posthog/schema.py b/posthog/schema.py index 5326b78126e65..dc77da163db17 100644 --- a/posthog/schema.py +++ b/posthog/schema.py @@ -4,7 +4,7 @@ from __future__ import annotations from enum import Enum -from typing import Any, Dict, List, Literal, Optional, TypeAlias, Union +from typing import Any, Dict, List, Literal, Optional, Union from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, RootModel @@ -3001,38 +3001,6 @@ class HogQLMetadata(BaseModel): table: Optional[str] = Field(default=None, description="Table to validate the expression against") -AnyQuery: TypeAlias = Union[ - EventsNode, - ActionsNode, - PersonsNode, - DataWarehouseNode, - TimeToSeeDataSessionsQuery, - EventsQuery, - ActorsQuery, - InsightActorsQuery, - InsightActorsQueryOptions, - SessionsTimelineQuery, - HogQLQuery, - HogQLMetadata, - HogQLAutocomplete, - WebOverviewQuery, - WebStatsTableQuery, - WebTopClicksQuery, - DataVisualizationNode, - DataTableNode, - SavedInsightNode, - InsightVizNode, - TrendsQuery, - FunnelsQuery, - RetentionQuery, - PathsQuery, - StickinessQuery, - LifecycleQuery, - FunnelCorrelationQuery, - DatabaseSchemaQuery, -] - - class QueryRequest(BaseModel): model_config = ConfigDict( extra="forbid", @@ -3046,7 +3014,36 @@ class QueryRequest(BaseModel): client_query_id: Optional[str] = Field( default=None, description="Client provided query ID. Can be used to retrieve the status or cancel the query." ) - query: AnyQuery = Field( + query: Union[ + EventsNode, + ActionsNode, + PersonsNode, + DataWarehouseNode, + TimeToSeeDataSessionsQuery, + EventsQuery, + ActorsQuery, + InsightActorsQuery, + InsightActorsQueryOptions, + SessionsTimelineQuery, + HogQLQuery, + HogQLMetadata, + HogQLAutocomplete, + WebOverviewQuery, + WebStatsTableQuery, + WebTopClicksQuery, + DataVisualizationNode, + DataTableNode, + SavedInsightNode, + InsightVizNode, + TrendsQuery, + FunnelsQuery, + RetentionQuery, + PathsQuery, + StickinessQuery, + LifecycleQuery, + FunnelCorrelationQuery, + DatabaseSchemaQuery, + ] = Field( ..., description='Submit a JSON string representing a query for PostHog data analysis, for example a HogQL query.\n\nExample payload:\n\n```\n\n{"query": {"kind": "HogQLQuery", "query": "select * from events limit 100"}}\n\n```\n\nFor more details on HogQL queries, see the [PostHog HogQL documentation](/docs/hogql#api-access).', discriminator="kind", @@ -3088,7 +3085,36 @@ class QuerySchemaRoot( ] ] ): - root: AnyQuery = Field(..., discriminator="kind") + root: Union[ + EventsNode, + ActionsNode, + PersonsNode, + DataWarehouseNode, + TimeToSeeDataSessionsQuery, + EventsQuery, + ActorsQuery, + InsightActorsQuery, + InsightActorsQueryOptions, + SessionsTimelineQuery, + HogQLQuery, + HogQLMetadata, + HogQLAutocomplete, + WebOverviewQuery, + WebStatsTableQuery, + WebTopClicksQuery, + DataVisualizationNode, + DataTableNode, + SavedInsightNode, + InsightVizNode, + TrendsQuery, + FunnelsQuery, + RetentionQuery, + PathsQuery, + StickinessQuery, + LifecycleQuery, + FunnelCorrelationQuery, + DatabaseSchemaQuery, + ] = Field(..., discriminator="kind") PropertyGroupFilterValue.model_rebuild()