Skip to content

Commit

Permalink
Add BaseScopedCount abstract model
Browse files Browse the repository at this point in the history
  • Loading branch information
rowanseymour committed Nov 28, 2024
1 parent 4584a44 commit 199410e
Show file tree
Hide file tree
Showing 5 changed files with 43 additions and 18 deletions.
10 changes: 3 additions & 7 deletions temba/flows/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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")
Expand Down
18 changes: 18 additions & 0 deletions temba/orgs/migrations/0160_alter_itemcount_scope.py
Original file line number Diff line number Diff line change
@@ -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),
),
]
10 changes: 3 additions & 7 deletions temba/orgs/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
14 changes: 13 additions & 1 deletion temba/utils/models/counts.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""
Expand All @@ -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
9 changes: 6 additions & 3 deletions temba/utils/models/squashable.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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:
Expand Down

0 comments on commit 199410e

Please sign in to comment.