diff --git a/posthog/hogql_queries/insights/test/test_retention_query_runner.py b/posthog/hogql_queries/insights/test/test_retention_query_runner.py index 30edb32102f76..04c108dd779f1 100644 --- a/posthog/hogql_queries/insights/test/test_retention_query_runner.py +++ b/posthog/hogql_queries/insights/test/test_retention_query_runner.py @@ -1,3 +1,5 @@ +from typing import Optional +from unittest.mock import MagicMock, patch import uuid from datetime import datetime @@ -6,11 +8,14 @@ from django.test import override_settings from rest_framework import status +from posthog.clickhouse.client.execute import sync_execute from posthog.constants import ( RETENTION_FIRST_TIME, TREND_FILTER_TYPE_ACTIONS, TREND_FILTER_TYPE_EVENTS, ) +from posthog.hogql.constants import LimitContext +from posthog.hogql.query import INCREASED_MAX_EXECUTION_TIME from posthog.hogql_queries.insights.retention_query_runner import RetentionQueryRunner from posthog.hogql_queries.actors_query_runner import ActorsQueryRunner from posthog.models import Action, ActionStep @@ -1685,10 +1690,10 @@ def test_day_interval_sampled(self): class TestClickhouseRetentionGroupAggregation(ClickhouseTestMixin, APIBaseTest): - def run_query(self, query): + def run_query(self, query, *, limit_context: Optional[LimitContext] = None): if not query.get("retentionFilter"): query["retentionFilter"] = {} - runner = RetentionQueryRunner(team=self.team, query=query) + runner = RetentionQueryRunner(team=self.team, query=query, limit_context=limit_context) return runner.calculate().model_dump()["results"] def run_actors_query(self, interval, query, select=None, actor="person"): @@ -1920,3 +1925,10 @@ def test_groups_aggregating_person_on_events(self): [1], ], ) + + @patch("posthog.hogql.query.sync_execute", wraps=sync_execute) + def test_limit_is_context_aware(self, mock_sync_execute: MagicMock): + self.run_query(query={}, limit_context=LimitContext.QUERY_ASYNC) + + mock_sync_execute.assert_called_once() + self.assertIn(f" max_execution_time={INCREASED_MAX_EXECUTION_TIME},", mock_sync_execute.call_args[0][0]) diff --git a/posthog/hogql_queries/insights/test/test_stickiness_query_runner.py b/posthog/hogql_queries/insights/test/test_stickiness_query_runner.py index 3de1fb6ce865e..6e25827e6ecba 100644 --- a/posthog/hogql_queries/insights/test/test_stickiness_query_runner.py +++ b/posthog/hogql_queries/insights/test/test_stickiness_query_runner.py @@ -1,8 +1,12 @@ from dataclasses import dataclass from typing import Dict, List, Optional, Union +from unittest.mock import MagicMock, patch from django.test import override_settings from freezegun import freeze_time +from posthog.clickhouse.client.execute import sync_execute +from posthog.hogql.constants import LimitContext +from posthog.hogql.query import INCREASED_MAX_EXECUTION_TIME from posthog.hogql_queries.insights.stickiness_query_runner import StickinessQueryRunner from posthog.models.action.action import Action from posthog.models.action_step import ActionStep @@ -197,6 +201,7 @@ def _run_query( properties: Optional[StickinessProperties] = None, filters: Optional[StickinessFilter] = None, filter_test_accounts: Optional[bool] = False, + limit_context: Optional[LimitContext] = None, ): query_series: List[EventsNode | ActionsNode] = [EventsNode(event="$pageview")] if series is None else series query_date_from = date_from or self.default_date_from @@ -211,7 +216,7 @@ def _run_query( stickinessFilter=filters, filterTestAccounts=filter_test_accounts, ) - return StickinessQueryRunner(team=self.team, query=query).calculate() + return StickinessQueryRunner(team=self.team, query=query, limit_context=limit_context).calculate() def test_stickiness_runs(self): self._create_test_events() @@ -580,3 +585,10 @@ def test_hogql_aggregations(self): 1, 0, ] + + @patch("posthog.hogql.query.sync_execute", wraps=sync_execute) + def test_limit_is_context_aware(self, mock_sync_execute: MagicMock): + self._run_query(limit_context=LimitContext.QUERY_ASYNC) + + mock_sync_execute.assert_called_once() + self.assertIn(f" max_execution_time={INCREASED_MAX_EXECUTION_TIME},", mock_sync_execute.call_args[0][0]) diff --git a/posthog/hogql_queries/insights/trends/test/test_trends_query_runner.py b/posthog/hogql_queries/insights/trends/test/test_trends_query_runner.py index 104e232a01406..1e10c5160ef89 100644 --- a/posthog/hogql_queries/insights/trends/test/test_trends_query_runner.py +++ b/posthog/hogql_queries/insights/trends/test/test_trends_query_runner.py @@ -1,11 +1,13 @@ from dataclasses import dataclass from typing import Dict, List, Optional -from unittest.mock import patch +from unittest.mock import MagicMock, patch from django.test import override_settings from freezegun import freeze_time +from posthog.clickhouse.client.execute import sync_execute from posthog.hogql import ast -from posthog.hogql.constants import MAX_SELECT_RETURNED_ROWS +from posthog.hogql.constants import MAX_SELECT_RETURNED_ROWS, LimitContext from posthog.hogql.modifiers import create_default_modifiers_for_team +from posthog.hogql.query import INCREASED_MAX_EXECUTION_TIME from posthog.hogql_queries.insights.trends.trends_query_runner import TrendsQueryRunner from posthog.models.cohort.cohort import Cohort from posthog.models.property_definition import PropertyDefinition @@ -175,6 +177,7 @@ def _create_query_runner( breakdown: Optional[BreakdownFilter] = None, filter_test_accounts: Optional[bool] = None, hogql_modifiers: Optional[HogQLQueryModifiers] = None, + limit_context: Optional[LimitContext] = None, ) -> TrendsQueryRunner: query_series: List[EventsNode | ActionsNode] = [EventsNode(event="$pageview")] if series is None else series query = TrendsQuery( @@ -185,7 +188,7 @@ def _create_query_runner( breakdownFilter=breakdown, filterTestAccounts=filter_test_accounts, ) - return TrendsQueryRunner(team=self.team, query=query, modifiers=hogql_modifiers) + return TrendsQueryRunner(team=self.team, query=query, modifiers=hogql_modifiers, limit_context=limit_context) def _run_trends_query( self, @@ -195,8 +198,10 @@ def _run_trends_query( series: Optional[List[EventsNode | ActionsNode]], trends_filters: Optional[TrendsFilter] = None, breakdown: Optional[BreakdownFilter] = None, + *, filter_test_accounts: Optional[bool] = None, hogql_modifiers: Optional[HogQLQueryModifiers] = None, + limit_context: Optional[LimitContext] = None, ): return self._create_query_runner( date_from=date_from, @@ -207,6 +212,7 @@ def _run_trends_query( breakdown=breakdown, filter_test_accounts=filter_test_accounts, hogql_modifiers=hogql_modifiers, + limit_context=limit_context, ).calculate() def test_trends_query_label(self): @@ -1649,3 +1655,16 @@ def test_to_actors_query_options_breakdowns_hogql(self): BreakdownItem(label="Safari", value="Safari"), BreakdownItem(label="Edge", value="Edge"), ] + + @patch("posthog.hogql.query.sync_execute", wraps=sync_execute) + def test_limit_is_context_aware(self, mock_sync_execute: MagicMock): + self._run_trends_query( + "2020-01-09", + "2020-01-20", + IntervalType.day, + [EventsNode(event="$pageview")], + limit_context=LimitContext.QUERY_ASYNC, + ) + + mock_sync_execute.assert_called_once() + self.assertIn(f" max_execution_time={INCREASED_MAX_EXECUTION_TIME},", mock_sync_execute.call_args[0][0]) diff --git a/posthog/hogql_queries/web_analytics/test/test_web_overview.py b/posthog/hogql_queries/web_analytics/test/test_web_overview.py index 63a26ffea9233..dcafe660fc72d 100644 --- a/posthog/hogql_queries/web_analytics/test/test_web_overview.py +++ b/posthog/hogql_queries/web_analytics/test/test_web_overview.py @@ -1,6 +1,11 @@ +from typing import Optional +from unittest.mock import MagicMock, patch from freezegun import freeze_time from parameterized import parameterized +from posthog.clickhouse.client.execute import sync_execute +from posthog.hogql.constants import LimitContext +from posthog.hogql.query import INCREASED_MAX_EXECUTION_TIME from posthog.hogql_queries.web_analytics.web_overview import WebOverviewQueryRunner from posthog.schema import WebOverviewQuery, DateRange from posthog.test.base import ( @@ -36,14 +41,21 @@ def _create_events(self, data, event="$pageview"): ) return person_result - def _run_web_overview_query(self, date_from, date_to, use_sessions_table=False, compare=True): + def _run_web_overview_query( + self, + date_from: str, + date_to: str, + use_sessions_table: bool = False, + compare: bool = True, + limit_context: Optional[LimitContext] = None, + ): query = WebOverviewQuery( dateRange=DateRange(date_from=date_from, date_to=date_to), properties=[], compare=compare, useSessionsTable=use_sessions_table, ) - runner = WebOverviewQueryRunner(team=self.team, query=query) + runner = WebOverviewQueryRunner(team=self.team, query=query, limit_context=limit_context) return runner.calculate() @parameterized.expand([(True,), (False,)]) @@ -185,3 +197,10 @@ def test_correctly_counts_pageviews_in_long_running_session(self, use_sessions_t sessions = results[2] self.assertEqual(1, sessions.value) + + @patch("posthog.hogql.query.sync_execute", wraps=sync_execute) + def test_limit_is_context_aware(self, mock_sync_execute: MagicMock): + self._run_web_overview_query("2023-12-01", "2023-12-03", limit_context=LimitContext.QUERY_ASYNC) + + mock_sync_execute.assert_called_once() + self.assertIn(f" max_execution_time={INCREASED_MAX_EXECUTION_TIME},", mock_sync_execute.call_args[0][0])