Read from new message folder counts

temba/flows/tests.py

@@ -21,7 +21,7 @@
 from temba.classifiers.models import Classifier
 from temba.contacts.models import URN, Contact, ContactField, ContactGroup, ContactURN
 from temba.globals.models import Global
-from temba.msgs.models import SystemLabel, SystemLabelCount
+from temba.msgs.models import SystemLabel
 from temba.orgs.integrations.dtone import DTOneType
 from temba.orgs.models import Export
 from temba.templates.models import TemplateTranslation
@@ -1933,7 +1933,7 @@ def test_preview_start(self, mr_mocks):
         )
 
         # if we have too many messages in our outbox we should block
-        SystemLabelCount.objects.create(org=self.org, label_type=SystemLabel.TYPE_OUTBOX, count=1_000_001)
+        self.org.counts.create(scope=f"msgs:folder:{SystemLabel.TYPE_OUTBOX}", count=1_000_001)
         preview_url = reverse("flows.flow_preview_start", args=[flow.id])
         mr_mocks.flow_start_preview(query="age > 30", total=1000)
 
@@ -1950,7 +1950,7 @@ def test_preview_start(self, mr_mocks):
             ],
             response.json()["blockers"],
         )
-        self.org.system_labels.all().delete()
+        self.org.counts.prefix("msgs:folder:").delete()
 
         # check warning for lots of contacts
         preview_url = reverse("flows.flow_preview_start", args=[flow.id])

temba/msgs/models.py

@@ -815,12 +815,9 @@ class SystemLabel:
     )
 
     @classmethod
-    def get_counts(cls, org, use_new=False):
-        if use_new:
-            counts = org.counts.prefix("msgs:folder:").scope_totals()
-            return {lb: counts.get(f"msgs:folder:{lb}", 0) for lb, n in cls.TYPE_CHOICES}
-
-        return SystemLabelCount.get_totals(org)
+    def get_counts(cls, org):
+        counts = org.counts.prefix("msgs:folder:").scope_totals()
+        return {lb: counts.get(f"msgs:folder:{lb}", 0) for lb, n in cls.TYPE_CHOICES}
 
     @classmethod
     def get_queryset(cls, org, label_type):
@@ -897,25 +894,14 @@ def get_archive_query(cls, label_type: str) -> dict:
 
 class SystemLabelCount(BaseSquashableCount):
     """
-    Counts of messages/broadcasts/calls maintained by database level triggers
+    TODO drop
     """
 
     squash_over = ("org_id", "label_type")
 
     org = models.ForeignKey(Org, on_delete=models.PROTECT, related_name="system_labels")
     label_type = models.CharField(max_length=1, choices=SystemLabel.TYPE_CHOICES)
 
-    @classmethod
-    def get_totals(cls, org):
-        """
-        Gets all system label counts by type for the given org
-        """
-        counts = cls.objects.filter(org=org).values_list("label_type").annotate(count_sum=Sum("count"))
-        counts_by_type = {c[0]: c[1] for c in counts}
-
-        # for convenience, include all label types
-        return {lb: counts_by_type.get(lb, 0) for lb, n in SystemLabel.TYPE_CHOICES}
-
     class Meta:
         indexes = [models.Index(fields=("org", "label_type", "is_squashed"))]
 

temba/msgs/tests.py

@@ -14,20 +14,10 @@
 from temba.archives.models import Archive
 from temba.channels.models import ChannelCount, ChannelLog
 from temba.flows.models import Flow
-from temba.msgs.models import (
-    Attachment,
-    Broadcast,
-    Label,
-    LabelCount,
-    Media,
-    MessageExport,
-    Msg,
-    OptIn,
-    SystemLabel,
-    SystemLabelCount,
-)
+from temba.msgs.models import Attachment, Broadcast, Label, LabelCount, Media, MessageExport, Msg, OptIn, SystemLabel
 from temba.msgs.views import ScheduleForm
 from temba.orgs.models import Export
+from temba.orgs.tasks import squash_item_counts
 from temba.schedules.models import Schedule
 from temba.templates.models import TemplateTranslation
 from temba.tests import CRUDLTestMixin, TembaTest, mock_mailroom, mock_uuids
@@ -2466,15 +2456,15 @@ def test_preview(self, mr_mocks):
 
         # if we have too many messages in our outbox we should block
         mr_mocks.msg_broadcast_preview(query="age > 30", total=2)
-        SystemLabelCount.objects.create(org=self.org, label_type=SystemLabel.TYPE_OUTBOX, count=1_000_001)
+        self.org.counts.create(scope=f"msgs:folder:{SystemLabel.TYPE_OUTBOX}", count=1_000_001)
         response = self.client.post(preview_url, {"query": "age > 30"}, content_type="application/json")
         self.assertEqual(
             [
                 "You have too many messages queued in your outbox. Please wait for these messages to send and then try again."
             ],
             response.json()["blockers"],
         )
-        self.org.system_labels.all().delete()
+        self.org.counts.prefix("msgs:folder:").delete()
 
         # if we release our send channel we can't send a broadcast
         self.channel.release(self.admin)
@@ -2899,8 +2889,7 @@ def test_get_archive_query(self):
 
     def test_get_counts(self):
         def assert_counts(org, expected: dict):
-            self.assertEqual(SystemLabel.get_counts(org, use_new=False), expected)
-            self.assertEqual(SystemLabel.get_counts(org, use_new=True), expected)
+            self.assertEqual(SystemLabel.get_counts(org), expected)
 
         assert_counts(
             self.org,
@@ -3021,10 +3010,10 @@ def assert_counts(org, expected: dict):
             },
         )
 
-        self.assertEqual(SystemLabelCount.objects.all().count(), 25)
+        self.assertEqual(self.org.counts.count(), 25)
 
         # squash our counts
-        squash_msg_counts()
+        squash_item_counts()
 
         assert_counts(
             self.org,
@@ -3040,8 +3029,8 @@ def assert_counts(org, expected: dict):
             },
         )
 
-        # we should only have one system label per type with no-zero count
-        self.assertEqual(SystemLabelCount.objects.all().count(), 5)
+        # we should only have one count per folder with non-zero count
+        self.assertEqual(self.org.counts.count(), 5)
 
 
 class TagsTest(TembaTest):
@@ -3188,7 +3177,7 @@ def setUpBeforeMigration(self, apps):
 
     def test_migration(self):
         def assert_counts(org, expected: dict):
-            self.assertEqual(SystemLabel.get_counts(org, use_new=True), expected)
+            self.assertEqual(SystemLabel.get_counts(org), expected)
 
         assert_counts(
             self.org,