perf: Use materialized columns in get_all_event_metrics_in_period (#2…

…5260)
PostHog · Sep 27, 2024 · abdf7a5 · abdf7a5
1 parent 6b6285d
commit abdf7a5
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 12 deletions.
diff --git a/posthog/tasks/test/test_usage_report.py b/posthog/tasks/test/test_usage_report.py
@@ -50,6 +50,7 @@
     ClickhouseTestMixin,
     _create_event,
     _create_person,
+    also_test_with_materialized_columns,
     flush_persons_and_events,
     snapshot_clickhouse_queries,
 )
@@ -813,6 +814,7 @@ def test_unlicensed_usage_report(self, mock_post: MagicMock, mock_client: MagicM
 
 @freeze_time("2022-01-09T00:01:00Z")
 class ReplayUsageReport(APIBaseTest, ClickhouseTestMixin, ClickhouseDestroyTablesMixin):
+    @also_test_with_materialized_columns(event_properties=["$lib"], verify_no_jsonextract=False)
     def test_usage_report_replay(self) -> None:
         _setup_replay_data(self.team.pk, include_mobile_replay=False)
 
@@ -831,6 +833,7 @@ def test_usage_report_replay(self) -> None:
         assert org_reports[str(self.organization.id)].recording_count_in_period == 5
         assert org_reports[str(self.organization.id)].mobile_recording_count_in_period == 0
 
+    @also_test_with_materialized_columns(event_properties=["$lib"], verify_no_jsonextract=False)
     def test_usage_report_replay_with_mobile(self) -> None:
         _setup_replay_data(self.team.pk, include_mobile_replay=True)
 
@@ -852,6 +855,7 @@ def test_usage_report_replay_with_mobile(self) -> None:
 
 
 class HogQLUsageReport(APIBaseTest, ClickhouseTestMixin, ClickhouseDestroyTablesMixin):
+    @also_test_with_materialized_columns(event_properties=["$lib"], verify_no_jsonextract=False)
     def test_usage_report_hogql_queries(self) -> None:
         for _ in range(0, 100):
             _create_event(

diff --git a/posthog/tasks/usage_report.py b/posthog/tasks/usage_report.py
@@ -17,6 +17,7 @@
 from retry import retry
 from sentry_sdk import capture_exception
 
+from ee.clickhouse.materialized_columns.columns import get_materialized_columns
 from posthog import version_requirement
 from posthog.clickhouse.client.connection import Workload
 from posthog.client import sync_execute
@@ -455,26 +456,31 @@ def get_teams_with_event_count_with_groups_in_period(begin: datetime, end: datet
 @timed_log()
 @retry(tries=QUERY_RETRIES, delay=QUERY_RETRY_DELAY, backoff=QUERY_RETRY_BACKOFF)
 def get_all_event_metrics_in_period(begin: datetime, end: datetime) -> dict[str, list[tuple[int, int]]]:
+    materialized_columns = get_materialized_columns("events")
+
+    # Check if $lib is materialized
+    lib_expression = materialized_columns.get(("$lib", "properties"), "JSONExtractString(properties, '$lib')")
+
     results = sync_execute(
-        """
+        f"""
         SELECT
             team_id,
             multiIf(
                 event LIKE 'helicone%%', 'helicone_events',
                 event LIKE 'langfuse%%', 'langfuse_events',
                 event LIKE 'keywords_ai%%', 'keywords_ai_events',
                 event LIKE 'traceloop%%', 'traceloop_events',
-                JSONExtractString(properties, '$lib') = 'web', 'web_events',
-                JSONExtractString(properties, '$lib') = 'posthog-node', 'node_events',
-                JSONExtractString(properties, '$lib') = 'posthog-android', 'android_events',
-                JSONExtractString(properties, '$lib') = 'posthog-flutter', 'flutter_events',
-                JSONExtractString(properties, '$lib') = 'posthog-ios', 'ios_events',
-                JSONExtractString(properties, '$lib') = 'posthog-go', 'go_events',
-                JSONExtractString(properties, '$lib') = 'posthog-java', 'java_events',
-                JSONExtractString(properties, '$lib') = 'posthog-react-native', 'react_native_events',
-                JSONExtractString(properties, '$lib') = 'posthog-ruby', 'ruby_events',
-                JSONExtractString(properties, '$lib') = 'posthog-python', 'python_events',
-                JSONExtractString(properties, '$lib') = 'posthog-php', 'php_events',
+                {lib_expression} = 'web', 'web_events',
+                {lib_expression} = 'posthog-node', 'node_events',
+                {lib_expression} = 'posthog-android', 'android_events',
+                {lib_expression} = 'posthog-flutter', 'flutter_events',
+                {lib_expression} = 'posthog-ios', 'ios_events',
+                {lib_expression} = 'posthog-go', 'go_events',
+                {lib_expression} = 'posthog-java', 'java_events',
+                {lib_expression} = 'posthog-react-native', 'react_native_events',
+                {lib_expression} = 'posthog-ruby', 'ruby_events',
+                {lib_expression} = 'posthog-python', 'python_events',
+                {lib_expression} = 'posthog-php', 'php_events',
                 'other'
             ) AS metric,
             count(1) as count