Add some tests

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])

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])

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])

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])