Merge pull request #5717 from nyaruka/base_count_model

Add `BaseScopedCount` abstract model

temba/flows/models.py

@@ -29,7 +29,7 @@
 from temba.utils import analytics, json, on_transaction_commit, s3
 from temba.utils.export.models import MultiSheetExporter
 from temba.utils.models import JSONAsTextField, LegacyUUIDMixin, SquashableModel, TembaModel, delete_in_batches
-from temba.utils.models.counts import ScopeCountQuerySet
+from temba.utils.models.counts import BaseScopedCount
 from temba.utils.uuid import uuid4
 
 from . import legacy
@@ -1426,18 +1426,14 @@ def release(self):
         self.delete()
 
 
-class FlowActivityCount(SquashableModel):
+class FlowActivityCount(BaseScopedCount):
     """
     Flow-level counts of activity.
     """
 
     squash_over = ("flow_id", "scope")
 
     flow = models.ForeignKey(Flow, on_delete=models.PROTECT, related_name="counts", db_index=False)  # indexed below
-    scope = models.CharField(max_length=128)
-    count = models.IntegerField(default=0)
-
-    objects = ScopeCountQuerySet.as_manager()
 
     @classmethod
     def get_squash_query(cls, distinct_set) -> tuple:
@@ -1460,7 +1456,7 @@ def get_squash_query(cls, distinct_set) -> tuple:
     @classmethod
     def prefetch_by_scope(cls, flows, *, prefix: str, to_attr: str, using: str):
         counts = (
-            FlowActivityCount.objects.using(using)
+            cls.objects.using(using)
             .filter(flow__in=flows)
             .prefix(prefix)
             .values_list("flow_id", "scope")

temba/orgs/migrations/0160_alter_itemcount_scope.py

@@ -0,0 +1,18 @@
+# Generated by Django 5.1.2 on 2024-11-28 19:50
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("orgs", "0159_usersettings_is_system"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="itemcount",
+            name="scope",
+            field=models.CharField(max_length=128),
+        ),
+    ]

temba/orgs/models.py

@@ -41,8 +41,8 @@
 from temba.utils.dates import datetime_to_str
 from temba.utils.email import EmailSender
 from temba.utils.fields import UploadToIdPathAndRename
-from temba.utils.models import JSONField, SquashableModel, TembaUUIDMixin, delete_in_batches
-from temba.utils.models.counts import ScopeCountQuerySet
+from temba.utils.models import JSONField, TembaUUIDMixin, delete_in_batches
+from temba.utils.models.counts import BaseScopedCount
 from temba.utils.s3 import public_file_storage
 from temba.utils.text import generate_secret, generate_token
 from temba.utils.timezones import timezone_to_country_code
@@ -1872,18 +1872,14 @@ def __repr__(self):  # pragma: no cover
         return f'<Export: id={self.id} type="{self.export_type}">'
 
 
-class ItemCount(SquashableModel):
+class ItemCount(BaseScopedCount):
     """
     Org-level counts of things.
     """
 
     squash_over = ("org_id", "scope")
 
     org = models.ForeignKey(Org, on_delete=models.PROTECT, related_name="counts", db_index=False)  # indexed below
-    scope = models.CharField(max_length=64)
-    count = models.IntegerField(default=0)
-
-    objects = ScopeCountQuerySet.as_manager()
 
     @classmethod
     def get_squash_query(cls, distinct_set) -> tuple:

temba/utils/models/counts.py

@@ -3,8 +3,10 @@
 
 from temba.utils.db.queries import or_list
 
+from .squashable import SquashableModel
 
-class ScopeCountQuerySet(models.QuerySet):
+
+class ScopedCountQuerySet(models.QuerySet):
     """
     Specialized queryset for scope + count models.
     """
@@ -30,3 +32,13 @@ def scope_totals(self) -> dict[str, int]:
         """
         counts = self.values_list("scope").annotate(count_sum=Sum("count"))
         return {c[0]: c[1] for c in counts}
+
+
+class BaseScopedCount(SquashableModel):
+    scope = models.CharField(max_length=128)
+    count = models.IntegerField(default=0)
+
+    objects = ScopedCountQuerySet.as_manager()
+
+    class Meta:
+        abstract = True

temba/utils/models/squashable.py

@@ -15,6 +15,10 @@ class SquashableModel(models.Model):
     id = models.BigAutoField(auto_created=True, primary_key=True)
     is_squashed = models.BooleanField(default=False)
 
+    @classmethod
+    def get_squash_over(cls) -> tuple:
+        return cls.squash_over
+
     @classmethod
     def get_unsquashed(cls):
         return cls.objects.filter(is_squashed=False)
@@ -27,9 +31,8 @@ def squash(cls) -> int:
         """
 
         num_sets = 0
-        distinct_sets = (
-            cls.get_unsquashed().order_by(*cls.squash_over).distinct(*cls.squash_over)[: cls.squash_max_distinct]
-        )
+        squash_over = cls.get_squash_over()
+        distinct_sets = cls.get_unsquashed().order_by(*squash_over).distinct(*squash_over)[: cls.squash_max_distinct]
 
         for distinct_set in distinct_sets:
             with connection.cursor() as cursor: