Rework msg exports to MsgFolder instead of SystemLabel

Author: rowanseymour

temba/contacts/models.py

@@ -644,10 +644,9 @@ def get_status_counts(cls, org) -> dict:
         return {g.group_type: count for g, count in ContactGroup.get_member_counts(groups).items()}
 
     def get_scheduled_broadcasts(self):
-        from temba.msgs.models import SystemLabel
-
         return (
-            SystemLabel.get_queryset(self.org, SystemLabel.TYPE_SCHEDULED)
+            self.org.broadcasts.filter(is_active=True)
+            .exclude(schedule=None)
             .filter(schedule__next_fire__gte=timezone.now())
             .filter(Q(contacts__in=[self]) | Q(groups__in=self.groups.all()))
             .select_related("org", "schedule")

temba/msgs/models.py

@@ -940,26 +940,6 @@ 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):
-        """
-        Gets the queryset for the given system label. Any change here needs to be reflected in a change to the db
-        trigger used to maintain the label counts.
-        """
-
-        assert label_type in [c[0] for c in cls.TYPE_CHOICES]
-
-        if label_type == cls.TYPE_SCHEDULED:
-            return org.broadcasts.filter(is_active=True).exclude(schedule=None)
-        elif label_type == cls.TYPE_CALLS:
-            return org.calls.all()
-
-        return MsgFolder.from_code(label_type).get_queryset(org)
-
-    @classmethod
-    def get_archive_query(cls, label_type: str) -> dict:
-        return MsgFolder.from_code(label_type).get_archive_query()
-
 
 class Label(TembaModel, DependencyMixin):
     """
@@ -1151,14 +1131,16 @@ def create(cls, org, user, start_date, end_date, system_label=None, label=None,
 
     def get_folder(self, export):
         label_uuid = export.config.get("label_uuid")
-        system_label = export.config.get("system_label")
+        folder_code = export.config.get("system_label")
         if label_uuid:
             return None, export.org.msgs_labels.filter(uuid=label_uuid).first()
+        elif folder_code:
+            return MsgFolder.from_code(folder_code), None
         else:
-            return system_label, None
+            return None, None
 
     def write(self, export):
-        system_label, label = self.get_folder(export)
+        folder, label = self.get_folder(export)
         start_date, end_date = export.get_date_range()
 
         # create our exporter
@@ -1172,7 +1154,7 @@ def write(self, export):
         num_records = 0
         logger.info(f"starting msgs export #{export.id} for org #{export.org.id}")
 
-        for batch in self._get_msg_batches(export, system_label, label, start_date, end_date):
+        for batch in self._get_msg_batches(export, folder, label, start_date, end_date):
             self._write_msgs(export, exporter, batch)
 
             num_records += len(batch)
@@ -1183,13 +1165,13 @@ def write(self, export):
 
         return *exporter.save_file(), num_records
 
-    def _get_msg_batches(self, export, system_label, label, start_date, end_date):
+    def _get_msg_batches(self, export, folder, label, start_date, end_date):
         from temba.archives.models import Archive
         from temba.flows.models import Flow
 
         # firstly get msgs from archives
-        if system_label:
-            where = SystemLabel.get_archive_query(system_label)
+        if folder:
+            where = folder.get_archive_query()
         elif label:
             where = {"visibility": "visible", "__raw__": f"'{label.uuid}' IN s.labels[*].uuid"}
         else:
@@ -1208,8 +1190,8 @@ def _get_msg_batches(self, export, system_label, label, start_date, end_date):
                 matching.append(record)
             yield matching
 
-        if system_label:
-            messages = SystemLabel.get_queryset(export.org, system_label)
+        if folder:
+            messages = folder.get_queryset(export.org)
         elif label:
             messages = label.get_messages()
         else:

temba/msgs/tests/test_systemlabel.py temba/msgs/tests/test_msg_folder.py

@@ -1,45 +1,45 @@
 from django.utils import timezone
 
 from temba.flows.models import Flow, FlowRun, FlowSession
-from temba.msgs.models import Msg, SystemLabel
+from temba.msgs.models import Msg, MsgFolder, SystemLabel
 from temba.orgs.tasks import squash_item_counts
 from temba.schedules.models import Schedule
 from temba.tests import TembaTest
 from temba.utils import s3
 
 
-class SystemLabelTest(TembaTest):
+class MsgFolderTest(TembaTest):
     def test_get_archive_query(self):
         tcs = (
             (
-                SystemLabel.TYPE_INBOX,
+                MsgFolder.INBOX,
                 "SELECT s.* FROM s3object s WHERE s.direction = 'in' AND s.visibility = 'visible' AND s.status = 'handled' AND s.flow IS NULL AND s.type != 'voice'",
             ),
             (
-                SystemLabel.TYPE_FLOWS,
+                MsgFolder.HANDLED,
                 "SELECT s.* FROM s3object s WHERE s.direction = 'in' AND s.visibility = 'visible' AND s.status = 'handled' AND s.flow IS NOT NULL AND s.type != 'voice'",
             ),
             (
-                SystemLabel.TYPE_ARCHIVED,
+                MsgFolder.ARCHIVED,
                 "SELECT s.* FROM s3object s WHERE s.direction = 'in' AND s.visibility = 'archived' AND s.status = 'handled' AND s.type != 'voice'",
             ),
             (
-                SystemLabel.TYPE_OUTBOX,
+                MsgFolder.OUTBOX,
                 "SELECT s.* FROM s3object s WHERE s.direction = 'out' AND s.visibility = 'visible' AND s.status IN ('initializing', 'queued', 'errored')",
             ),
             (
-                SystemLabel.TYPE_SENT,
+                MsgFolder.SENT,
                 "SELECT s.* FROM s3object s WHERE s.direction = 'out' AND s.visibility = 'visible' AND s.status IN ('wired', 'sent', 'delivered', 'read')",
             ),
             (
-                SystemLabel.TYPE_FAILED,
+                MsgFolder.FAILED,
                 "SELECT s.* FROM s3object s WHERE s.direction = 'out' AND s.visibility = 'visible' AND s.status = 'failed'",
             ),
         )
 
-        for label_type, expected_select in tcs:
-            select = s3.compile_select(where=SystemLabel.get_archive_query(label_type))
-            self.assertEqual(expected_select, select, f"select s3 mismatch for label {label_type}")
+        for folder, expected_select in tcs:
+            select = s3.compile_select(where=folder.get_archive_query())
+            self.assertEqual(expected_select, select, f"select s3 mismatch for {folder}")
 
     def test_get_counts(self):
         def assert_counts(org, expected: dict):

temba/msgs/views.py

@@ -44,7 +44,7 @@
 from temba.utils.views.wizard import SmartWizardUpdateView, SmartWizardView
 
 from .forms import ComposeForm, ScheduleForm, TargetForm
-from .models import Broadcast, Label, LabelCount, Media, MessageExport, Msg, OptIn, SystemLabel
+from .models import Broadcast, Label, LabelCount, Media, MessageExport, Msg, MsgFolder, OptIn, SystemLabel
 
 
 class MsgListView(ContextMenuMixin, BulkActionMixin, SpaMixin, BaseListView):
@@ -58,20 +58,20 @@ class MsgListView(ContextMenuMixin, BulkActionMixin, SpaMixin, BaseListView):
     allow_export = False
     bulk_actions = ()
     bulk_action_permissions = {"resend": "msgs.msg_create", "delete": "msgs.msg_update"}
-    system_label = None
+    folder = None
     paginate_by = 100
 
     def pre_process(self, request, *args, **kwargs):
-        if self.system_label:
-            self.queryset = SystemLabel.get_queryset(request.org, self.system_label)
+        if self.folder:
+            self.queryset = self.folder.get_queryset(request.org)
 
-    def derive_label(self):
-        return self.system_label
+    def derive_folder(self):
+        return self.folder
 
     def derive_export_url(self):
         redirect = quote_plus(self.request.get_full_path())
-        label = self.derive_label()
-        label_id = label.uuid if isinstance(label, Label) else label
+        folder = self.derive_folder()
+        label_id = folder.uuid if isinstance(folder, Label) else folder.code
         return "%s?l=%s&redirect=%s" % (reverse("msgs.msg_export"), label_id, redirect)
 
     def get_queryset(self, **kwargs):
@@ -91,14 +91,14 @@ def get_queryset(self, **kwargs):
     def get_context_data(self, **kwargs):
         org = self.request.org
         counts = SystemLabel.get_counts(org)
-        label = self.derive_label()
+        folder = self.derive_folder()
 
         # if there isn't a search filtering the queryset, we can replace the count function with a pre-calculated value
         if "search" not in self.request.GET:
-            if isinstance(label, Label):
-                patch_queryset_count(self.object_list, label.get_visible_count)
-            elif isinstance(label, str):
-                patch_queryset_count(self.object_list, lambda: counts[label])
+            if isinstance(folder, Label):
+                patch_queryset_count(self.object_list, folder.get_visible_count)
+            elif isinstance(folder, MsgFolder):
+                patch_queryset_count(self.object_list, lambda: counts[folder.code])
 
         context = super().get_context_data(**kwargs)
         context["has_messages"] = (
@@ -699,7 +699,7 @@ def derive_url_pattern(cls, path, action):
     class Inbox(MsgListView):
         title = _("Inbox")
         template_name = "msgs/message_box.html"
-        system_label = SystemLabel.TYPE_INBOX
+        folder = MsgFolder.INBOX
         bulk_actions = ("archive", "label")
         allow_export = True
         menu_path = "/msg/inbox"
@@ -715,7 +715,7 @@ def get_queryset(self, **kwargs):
     class Flow(MsgListView):
         title = _("Handled")
         template_name = "msgs/message_box.html"
-        system_label = SystemLabel.TYPE_FLOWS
+        folder = MsgFolder.HANDLED
         bulk_actions = ("archive", "label")
         allow_export = True
         menu_path = "/msg/handled"
@@ -727,7 +727,7 @@ def get_queryset(self, **kwargs):
     class Archived(MsgListView):
         title = _("Archived")
         template_name = "msgs/message_box.html"
-        system_label = SystemLabel.TYPE_ARCHIVED
+        folder = MsgFolder.ARCHIVED
         bulk_actions = ("restore", "label", "delete")
         allow_export = True
 
@@ -738,7 +738,7 @@ def get_queryset(self, **kwargs):
     class Outbox(MsgListView):
         title = _("Outbox")
         template_name = "msgs/message_box.html"
-        system_label = SystemLabel.TYPE_OUTBOX
+        folder = MsgFolder.OUTBOX
         bulk_actions = ()
         allow_export = True
 
@@ -748,7 +748,7 @@ def get_queryset(self, **kwargs):
     class Sent(MsgListView):
         title = _("Sent")
         template_name = "msgs/msg_sent.html"
-        system_label = SystemLabel.TYPE_SENT
+        folder = MsgFolder.SENT
         bulk_actions = ()
         allow_export = True
         default_order = ("-sent_on", "-id")
@@ -759,7 +759,7 @@ def get_queryset(self, **kwargs):
     class Failed(MsgListView):
         title = _("Failed")
         template_name = "msgs/message_box.html"
-        system_label = SystemLabel.TYPE_FAILED
+        folder = MsgFolder.FAILED
         allow_export = True
 
         def get_bulk_actions(self):
@@ -810,7 +810,7 @@ def derive_url_pattern(cls, path, action):
         def label(self):
             return self.request.org.msgs_labels.get(uuid=self.kwargs["label_uuid"])
 
-        def derive_label(self):
+        def derive_folder(self):
             return self.label
 
         def get_queryset(self, **kwargs):