From 199410e7f2458f08c56f16f7b3d2caabfe240889 Mon Sep 17 00:00:00 2001 From: Rowan Seymour Date: Thu, 28 Nov 2024 19:52:46 +0000 Subject: [PATCH] Add BaseScopedCount abstract model --- temba/flows/models.py | 10 +++------- .../migrations/0160_alter_itemcount_scope.py | 18 ++++++++++++++++++ temba/orgs/models.py | 10 +++------- temba/utils/models/counts.py | 14 +++++++++++++- temba/utils/models/squashable.py | 9 ++++++--- 5 files changed, 43 insertions(+), 18 deletions(-) create mode 100644 temba/orgs/migrations/0160_alter_itemcount_scope.py diff --git a/temba/flows/models.py b/temba/flows/models.py index fd54af1881..16f8c85dea 100644 --- a/temba/flows/models.py +++ b/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,7 +1426,7 @@ def release(self): self.delete() -class FlowActivityCount(SquashableModel): +class FlowActivityCount(BaseScopedCount): """ Flow-level counts of activity. """ @@ -1434,10 +1434,6 @@ class FlowActivityCount(SquashableModel): 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") diff --git a/temba/orgs/migrations/0160_alter_itemcount_scope.py b/temba/orgs/migrations/0160_alter_itemcount_scope.py new file mode 100644 index 0000000000..79de97b8c6 --- /dev/null +++ b/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), + ), + ] diff --git a/temba/orgs/models.py b/temba/orgs/models.py index 9d4104ef61..593efade59 100644 --- a/temba/orgs/models.py +++ b/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,7 +1872,7 @@ def __repr__(self): # pragma: no cover return f'' -class ItemCount(SquashableModel): +class ItemCount(BaseScopedCount): """ Org-level counts of things. """ @@ -1880,10 +1880,6 @@ class ItemCount(SquashableModel): 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: diff --git a/temba/utils/models/counts.py b/temba/utils/models/counts.py index 52651c9719..9167270aa6 100644 --- a/temba/utils/models/counts.py +++ b/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 diff --git a/temba/utils/models/squashable.py b/temba/utils/models/squashable.py index 171f1060e7..91dce580cc 100644 --- a/temba/utils/models/squashable.py +++ b/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: