fix(insights): Pass Limit Context to breakdowns (#22169)

PostHog · May 7, 2024 · bb1e0aa · bb1e0aa
1 parent f5342b7
commit bb1e0aa
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/posthog/hogql_queries/insights/trends/breakdown_values.py b/posthog/hogql_queries/insights/trends/breakdown_values.py
@@ -75,6 +75,7 @@ def __init__(
         self.breakdown_limit = breakdown_filter.breakdown_limit or get_breakdown_limit_for_context(limit_context)
         self.query_date_range = query_date_range
         self.modifiers = modifiers
+        self.limit_context = limit_context
 
     def get_breakdown_values(self) -> list[str | int]:
         if self.breakdown_type == "cohort":
@@ -202,6 +203,7 @@ def get_breakdown_values(self) -> list[str | int]:
                 query=query,
                 team=self.team,
                 modifiers=self.modifiers,
+                limit_context=self.limit_context,
             )
             if response.results and len(response.results) > 0:
                 values = response.results[0][0]
@@ -215,6 +217,7 @@ def get_breakdown_values(self) -> list[str | int]:
                 query=query,
                 team=self.team,
                 modifiers=self.modifiers,
+                limit_context=self.limit_context,
             )
             value_index = (response.columns or []).index("value")
             values = [row[value_index] for row in response.results or []]

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
@@ -658,6 +658,24 @@ def test_formula_with_breakdown_and_no_data(self):
         )
         self.assertEqual([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], response.results[0]["data"])
 
+    @patch("posthog.hogql.query.sync_execute", wraps=sync_execute)
+    def test_breakdown_is_context_aware(self, mock_sync_execute: MagicMock):
+        self._create_test_events()
+
+        self._run_trends_query(
+            self.default_date_from,
+            self.default_date_to,
+            IntervalType.day,
+            [EventsNode(event="$pageviewxxx"), EventsNode(event="$pageleavexxx")],
+            TrendsFilter(formula="A+2*B"),
+            BreakdownFilter(breakdown_type=BreakdownType.person, breakdown="$browser"),
+            limit_context=LimitContext.QUERY_ASYNC,
+        )
+
+        self.assertEqual(mock_sync_execute.call_count, 4)
+        for mock_execute_call_args in mock_sync_execute.call_args_list:
+            self.assertIn(f" max_execution_time={INCREASED_MAX_EXECUTION_TIME},", mock_execute_call_args[0][0])
+
     def test_trends_compare(self):
         self._create_test_events()