From 1ff14fd0f1b0ca495c8508d9015544ac6debc24d Mon Sep 17 00:00:00 2001
From: Tim Abbott <tabbott@zulip.com>
Date: Wed, 2 Oct 2024 10:40:59 -0700
Subject: [PATCH] analytics: Pass subgroup=None to improve indexing.

Because the main indexes on end_time either don't include realm_id or
do include subgroup, passing an explicit subgroup=None for
single-realm queries to read CountStats that don't use the subgroups
feature greatly improves the query plans.
---
 corporate/lib/support.py    | 4 ++++
 zerver/actions/invites.py   | 3 +++
 zerver/lib/cache_helpers.py | 8 +++++++-
 zerver/models/realms.py     | 4 +++-
 4 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/corporate/lib/support.py b/corporate/lib/support.py
index 992c5ba760..23cc5933a8 100644
--- a/corporate/lib/support.py
+++ b/corporate/lib/support.py
@@ -346,10 +346,12 @@ def get_mobile_push_data(remote_entity: RemoteZulipServer | RemoteRealm) -> Mobi
         mobile_pushes = RemoteInstallationCount.objects.filter(
             server=remote_entity,
             property="mobile_pushes_forwarded::day",
+            subgroup=None,
             end_time__gte=timezone_now() - timedelta(days=7),
         ).aggregate(total_forwarded=Sum("value", default=0))
         latest_remote_server_push_forwarded_count = RemoteInstallationCount.objects.filter(
             server=remote_entity,
+            subgroup=None,
             property="mobile_pushes_forwarded::day",
         ).last()
         if latest_remote_server_push_forwarded_count is not None:  # nocoverage
@@ -380,10 +382,12 @@ def get_mobile_push_data(remote_entity: RemoteZulipServer | RemoteRealm) -> Mobi
         mobile_pushes = RemoteRealmCount.objects.filter(
             remote_realm=remote_entity,
             property="mobile_pushes_forwarded::day",
+            subgroup=None,
             end_time__gte=timezone_now() - timedelta(days=7),
         ).aggregate(total_forwarded=Sum("value", default=0))
         latest_remote_realm_push_forwarded_count = RemoteRealmCount.objects.filter(
             remote_realm=remote_entity,
+            subgroup=None,
             property="mobile_pushes_forwarded::day",
         ).last()
         if latest_remote_realm_push_forwarded_count is not None:  # nocoverage
diff --git a/zerver/actions/invites.py b/zerver/actions/invites.py
index c52a71c47e..2197ac0da3 100644
--- a/zerver/actions/invites.py
+++ b/zerver/actions/invites.py
@@ -41,6 +41,9 @@ def estimate_recent_invites(realms: Collection[Realm] | QuerySet[Realm], *, days
     recent_invites = RealmCount.objects.filter(
         realm__in=realms,
         property="invites_sent::day",
+        # It's important to filter on this even though the count
+        # doesn't use subgroup, so that we use the index.
+        subgroup=None,
         end_time__gte=timezone_now() - timedelta(days=days),
     ).aggregate(Sum("value"))["value__sum"]
     if recent_invites is None:
diff --git a/zerver/lib/cache_helpers.py b/zerver/lib/cache_helpers.py
index 9d940bd0d3..4494f4bfc2 100644
--- a/zerver/lib/cache_helpers.py
+++ b/zerver/lib/cache_helpers.py
@@ -67,7 +67,13 @@ def get_active_realm_ids() -> QuerySet[RealmCount, int]:
     """
     date = timezone_now() - timedelta(days=2)
     return (
-        RealmCount.objects.filter(end_time__gte=date, property="1day_actives::day", value__gt=0)
+        RealmCount.objects.filter(
+            end_time__gte=date,
+            property="1day_actives::day",
+            # Filtering on subgroup is important to ensure we use the good indexes.
+            subgroup=None,
+            value__gt=0,
+        )
         .distinct("realm_id")
         .values_list("realm_id", flat=True)
     )
diff --git a/zerver/models/realms.py b/zerver/models/realms.py
index ec4a766714..6ba3177d93 100644
--- a/zerver/models/realms.py
+++ b/zerver/models/realms.py
@@ -1022,7 +1022,9 @@ class Realm(models.Model):  # type: ignore[django-manager-missing] # django-stub
 
         try:
             latest_count_stat = RealmCount.objects.filter(
-                realm=realm, property="upload_quota_used_bytes::day"
+                realm=realm,
+                property="upload_quota_used_bytes::day",
+                subgroup=None,
             ).latest("end_time")
             last_recorded_used_space = latest_count_stat.value
             last_recorded_date = latest_count_stat.end_time