From 07fc6ffeedd2c627afa53adedb45691b8c30435c Mon Sep 17 00:00:00 2001
From: Tom Owers <tom@paace.co>
Date: Fri, 13 Oct 2023 15:29:32 +0100
Subject: [PATCH] Add timings measure for breakdown values query

---
 .../insights/trends/breakdown.py              | 31 ++++++++++++-------
 .../insights/trends/query_builder.py          | 10 +++++-
 .../insights/trends/trends_query_runner.py    |  2 +-
 3 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/posthog/hogql_queries/insights/trends/breakdown.py b/posthog/hogql_queries/insights/trends/breakdown.py
index d216a928afdd7..036675cc62716 100644
--- a/posthog/hogql_queries/insights/trends/breakdown.py
+++ b/posthog/hogql_queries/insights/trends/breakdown.py
@@ -1,5 +1,6 @@
 from typing import Dict, List, Tuple
 from posthog.hogql import ast
+from posthog.hogql.timings import HogQLTimings
 from posthog.hogql_queries.insights.trends.breakdown_values import BreakdownValues
 from posthog.hogql_queries.insights.trends.utils import get_properties_chain, series_event_name
 from posthog.hogql_queries.utils.query_date_range import QueryDateRange
@@ -13,14 +14,21 @@ class Breakdown:
     team: Team
     series: EventsNode | ActionsNode
     query_date_range: QueryDateRange
+    timings: HogQLTimings
 
     def __init__(
-        self, team: Team, query: TrendsQuery, series: EventsNode | ActionsNode, query_date_range: QueryDateRange
+        self,
+        team: Team,
+        query: TrendsQuery,
+        series: EventsNode | ActionsNode,
+        query_date_range: QueryDateRange,
+        timings: HogQLTimings,
     ):
         self.team = team
         self.query = query
         self.series = series
         self.query_date_range = query_date_range
+        self.timings = timings
 
     @cached_property
     def enabled(self) -> bool:
@@ -65,16 +73,17 @@ def _breakdown_values_ast(self) -> ast.Array:
 
     @cached_property
     def _get_breakdown_values(self) -> ast.Array:
-        breakdown = BreakdownValues(
-            team=self.team,
-            event_name=series_event_name(self.series),
-            breakdown_field=self.query.breakdown.breakdown,
-            breakdown_type=self.query.breakdown.breakdown_type,
-            query_date_range=self.query_date_range,
-            histogram_bin_count=self.query.breakdown.breakdown_histogram_bin_count,
-            group_type_index=self.query.breakdown.breakdown_group_type_index,
-        )
-        return breakdown.get_breakdown_values()
+        with self.timings.measure("breakdown_values_query"):
+            breakdown = BreakdownValues(
+                team=self.team,
+                event_name=series_event_name(self.series),
+                breakdown_field=self.query.breakdown.breakdown,
+                breakdown_type=self.query.breakdown.breakdown_type,
+                query_date_range=self.query_date_range,
+                histogram_bin_count=self.query.breakdown.breakdown_histogram_bin_count,
+                group_type_index=self.query.breakdown.breakdown_group_type_index,
+            )
+            return breakdown.get_breakdown_values()
 
     def _get_breakdown_histogram_buckets(self) -> List[Tuple[float, float]]:
         buckets = []
diff --git a/posthog/hogql_queries/insights/trends/query_builder.py b/posthog/hogql_queries/insights/trends/query_builder.py
index a7b4f3f401fb5..efb3e349a3333 100644
--- a/posthog/hogql_queries/insights/trends/query_builder.py
+++ b/posthog/hogql_queries/insights/trends/query_builder.py
@@ -2,6 +2,7 @@
 from posthog.hogql import ast
 from posthog.hogql.parser import parse_expr, parse_select
 from posthog.hogql.property import property_to_expr
+from posthog.hogql.timings import HogQLTimings
 from posthog.hogql_queries.insights.trends.breakdown import Breakdown
 from posthog.hogql_queries.insights.trends.utils import series_event_name
 from posthog.hogql_queries.utils.query_date_range import QueryDateRange
@@ -15,14 +16,21 @@ class TrendsQueryBuilder:
     team: Team
     query_date_range: QueryDateRange
     series: EventsNode | ActionsNode
+    timings: HogQLTimings
 
     def __init__(
-        self, trends_query: TrendsQuery, team: Team, query_date_range: QueryDateRange, series: EventsNode | ActionsNode
+        self,
+        trends_query: TrendsQuery,
+        team: Team,
+        query_date_range: QueryDateRange,
+        series: EventsNode | ActionsNode,
+        timings: HogQLTimings,
     ):
         self.query = trends_query
         self.team = team
         self.query_date_range = query_date_range
         self.series = series
+        self.timings = timings
 
     def build_query(self) -> ast.SelectUnionQuery:
         date_subqueries = self._get_date_subqueries()
diff --git a/posthog/hogql_queries/insights/trends/trends_query_runner.py b/posthog/hogql_queries/insights/trends/trends_query_runner.py
index 4e8f556d1a6df..840180fd0d99a 100644
--- a/posthog/hogql_queries/insights/trends/trends_query_runner.py
+++ b/posthog/hogql_queries/insights/trends/trends_query_runner.py
@@ -63,7 +63,7 @@ def to_query(self) -> List[ast.SelectQuery]:
                 else:
                     query_date_range = self.query_previous_date_range
 
-                query_builder = TrendsQueryBuilder(self.query, self.team, query_date_range, series.series)
+                query_builder = TrendsQueryBuilder(self.query, self.team, query_date_range, series.series, self.timings)
                 queries.append(query_builder.build_query())
 
         return queries