Merge pull request #5612 from nyaruka/notification_counts_v2_pt2

New notification counts part 2

temba/notifications/migrations/0026_backfill_new_counts.py

@@ -0,0 +1,26 @@
+# Generated by Django 5.1.2 on 2024-11-05 16:22
+
+from django.db import migrations, transaction
+
+
+def backfill_item_counts(apps, schema_editor):  # pragma: no cover
+    Org = apps.get_model("orgs", "Org")
+
+    for org in Org.objects.all():
+        user_ids = org.notifications.values_list("user_id", flat=True).distinct()
+
+        for user_id in user_ids:
+            with transaction.atomic():
+                seen_count = org.notifications.filter(user_id=user_id, medium__contains="U", is_seen=True).count()
+                unseen_count = org.notifications.filter(user_id=user_id, medium__contains="U", is_seen=False).count()
+
+                org.counts.filter(scope__startswith=f"notifications:{user_id}:").delete()
+                org.counts.create(scope=f"notifications:{user_id}:S", count=seen_count, is_squashed=True)
+                org.counts.create(scope=f"notifications:{user_id}:U", count=unseen_count, is_squashed=True)
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [("notifications", "0025_update_triggers")]
+
+    operations = [migrations.RunPython(backfill_item_counts, migrations.RunPython.noop)]

temba/notifications/models.py

@@ -224,10 +224,6 @@ def mark_seen(cls, org, user, notification_type: str = None, *, scope: str = Non
 
     @classmethod
     def get_unseen_count(cls, org: Org, user: User) -> int:
-        return NotificationCount.get_total(org, user)
-
-    @classmethod
-    def get_unseen_count_new(cls, org: Org, user: User) -> int:
         return org.counts.filter(scope=f"notifications:{user.id}:U").sum()
 
     @property
@@ -269,29 +265,11 @@ class Meta:
 
 class NotificationCount(SquashableModel):
     """
-    A count of a user's unseen notifications in a specific org
+    TODO drop
     """
 
     squash_over = ("org_id", "user_id")
 
     org = models.ForeignKey(Org, on_delete=models.PROTECT, related_name="notification_counts")
     user = models.ForeignKey(User, on_delete=models.PROTECT, related_name="notification_counts")
     count = models.IntegerField(default=0)
-
-    @classmethod
-    def get_squash_query(cls, distinct_set):
-        sql = """
-            WITH deleted as (
-                DELETE FROM %(table)s WHERE "org_id" = %%s AND "user_id" = %%s RETURNING "count"
-            )
-            INSERT INTO %(table)s("org_id", "user_id", "count", "is_squashed")
-            VALUES (%%s, %%s, GREATEST(0, (SELECT SUM("count") FROM deleted)), TRUE);
-            """ % {
-            "table": cls._meta.db_table
-        }
-
-        return sql, (distinct_set.org_id, distinct_set.user_id) * 2
-
-    @classmethod
-    def get_total(cls, org: Org, user: User) -> int:
-        return cls.sum(cls.objects.filter(org=org, user=user))

temba/notifications/tasks.py

@@ -6,7 +6,7 @@
 from temba.utils.crons import cron_task
 from temba.utils.models import delete_in_batches
 
-from .models import Notification, NotificationCount
+from .models import Notification
 
 logger = logging.getLogger(__name__)
 
@@ -32,11 +32,6 @@ def send_notification_emails():
     return {"sent": num_sent, "errored": num_errored}
 
 
-@cron_task(lock_timeout=1800)
-def squash_notification_counts():
-    NotificationCount.squash()
-
-
 @cron_task()
 def trim_notifications():
     trim_before = timezone.now() - settings.RETENTION_PERIODS["notification"]

temba/notifications/tests.py

@@ -7,13 +7,13 @@
 from temba.contacts.models import ContactExport, ContactImport
 from temba.flows.models import ResultsExport
 from temba.msgs.models import MessageExport
-from temba.orgs.models import Invitation, OrgRole
-from temba.tests import CRUDLTestMixin, TembaTest, matchers
+from temba.orgs.models import Invitation, ItemCount, OrgRole
+from temba.tests import CRUDLTestMixin, MigrationTest, TembaTest, matchers
 from temba.tickets.models import TicketExport
 
 from .incidents.builtin import ChannelTemplatesFailedIncidentType, OrgFlaggedIncidentType
 from .models import Incident, Notification
-from .tasks import send_notification_emails, squash_notification_counts, trim_notifications
+from .tasks import send_notification_emails, trim_notifications
 from .types.builtin import (
     ExportFinishedNotificationType,
     InvitationAcceptedNotificationType,
@@ -545,7 +545,6 @@ def test_counts(self):
 
         def assert_count(org, user, expected: int):
             self.assertEqual(expected, Notification.get_unseen_count(org, user))
-            self.assertEqual(expected, Notification.get_unseen_count_new(org, user))
 
         assert_count(self.org, self.agent, 2)
         assert_count(self.org, self.editor, 3)
@@ -566,7 +565,7 @@ def assert_count(org, user, expected: int):
         assert_count(self.org2, self.agent, 0)
         assert_count(self.org2, self.editor, 1)
 
-        squash_notification_counts()
+        ItemCount.squash()
 
         assert_count(self.org, self.agent, 1)
         assert_count(self.org, self.editor, 2)
@@ -598,3 +597,34 @@ def test_trim_task(self):
 
         self.assertFalse(Notification.objects.filter(id=notification1.id).exists())
         self.assertTrue(Notification.objects.filter(id=notification2.id).exists())
+
+
+class BackfillNewCountsTest(MigrationTest):
+    app = "notifications"
+    migrate_from = "0025_update_triggers"
+    migrate_to = "0026_backfill_new_counts"
+
+    def setUpBeforeMigration(self, apps):
+        imp = ContactImport.objects.create(
+            org=self.org, mappings={}, num_records=5, created_by=self.editor, modified_by=self.editor
+        )
+        Notification.create_all(
+            imp.org, "import:finished", scope=f"contact:{imp.id}", users=[self.editor], contact_import=imp, medium="UE"
+        )
+        Notification.create_all(self.org, "tickets:opened", scope="", users=[self.agent, self.editor], medium="UE")
+        Notification.create_all(self.org, "tickets:activity", scope="", users=[self.agent, self.editor], medium="UE")
+        Notification.create_all(self.org, "tickets:reply", scope="12", users=[self.editor], medium="E")  # email only
+        Notification.create_all(
+            self.org2, "tickets:activity", scope="", users=[self.editor], medium="UE"
+        )  # different org
+
+        self.org2.counts.all().delete()
+
+    def test_migration(self):
+        def assert_count(org, user, expected: int):
+            self.assertEqual(expected, Notification.get_unseen_count(org, user))
+
+        assert_count(self.org, self.agent, 2)
+        assert_count(self.org, self.editor, 3)
+        assert_count(self.org2, self.agent, 0)
+        assert_count(self.org2, self.editor, 1)

temba/settings_common.py

@@ -764,9 +764,8 @@
     "squash-channel-counts": {"task": "squash_channel_counts", "schedule": timedelta(seconds=60)},
     "squash-group-counts": {"task": "squash_group_counts", "schedule": timedelta(seconds=60)},
     "squash-flow-counts": {"task": "squash_flow_counts", "schedule": timedelta(seconds=60)},
-    "squash-item-counts": {"task": "squash_item_counts", "schedule": timedelta(seconds=45)},
+    "squash-item-counts": {"task": "squash_item_counts", "schedule": timedelta(seconds=30)},
     "squash-msg-counts": {"task": "squash_msg_counts", "schedule": timedelta(seconds=60)},
-    "squash-notification-counts": {"task": "squash_notification_counts", "schedule": timedelta(seconds=60)},
     "squash-ticket-counts": {"task": "squash_ticket_counts", "schedule": timedelta(seconds=60)},
     "sync-classifier-intents": {"task": "sync_classifier_intents", "schedule": timedelta(seconds=300)},
     "track-org-channel-counts": {"task": "track_org_channel_counts", "schedule": crontab(hour=4, minute=0)},