Move more counts into with_event_counts

djstripe_ext/tasks.py

@@ -1,57 +1,35 @@
 from django.conf import settings
-from django.db.models import Count, Q, F, Subquery, OuterRef, IntegerField
+from django.db.models import Q, F, IntegerField
 from django.db.models.functions import Cast
 from django.db.models.fields.json import KeyTextTransform
 from celery import shared_task
 from organizations_ext.models import Organization
-from projects.models import Project
 from .models import SubscriptionQuotaWarning
 from .email import WarnQuotaEmail
 
 
 @shared_task
 def warn_organization_throttle():
-    """ Warn user about approaching 80% of allotted events """
+    """Warn user about approaching 80% of allotted events"""
     if not settings.BILLING_ENABLED:
         return
 
-    queryset = Organization.objects.filter(
-        djstripe_customers__subscriptions__status="active",
-    ).filter(
-        Q(djstripe_customers__subscriptions__subscriptionquotawarning=None)
-        | Q(
-            djstripe_customers__subscriptions__subscriptionquotawarning__notice_last_sent__lt=F(
-                "djstripe_customers__subscriptions__current_period_start"
-            ),
+    queryset = (
+        Organization.objects.with_event_counts()
+        .filter(
+            djstripe_customers__subscriptions__status="active",
         )
-    )
-
-    projects = Project.objects.filter(organization=OuterRef("pk")).values(
-        "organization"
-    )
-    total_issue_events = projects.annotate(
-        total=Count(
-            "issue__event",
-            filter=Q(
-                issue__event__created__gte=OuterRef(
+        .filter(
+            Q(djstripe_customers__subscriptions__subscriptionquotawarning=None)
+            | Q(
+                djstripe_customers__subscriptions__subscriptionquotawarning__notice_last_sent__lt=F(
                     "djstripe_customers__subscriptions__current_period_start"
-                )
-            ),
+                ),
+            )
         )
-    ).values("total")
-    total_transaction_events = projects.annotate(
-        total=Count(
-            "transactiongroup__transactionevent",
-            filter=Q(
-                transactiongroup__transactionevent__created__gte=OuterRef(
-                    "djstripe_customers__subscriptions__current_period_start"
-                )
-            ),
-        )
-    ).values("total")
+    )
 
     queryset = queryset.annotate(
-        event_count=Subquery(total_issue_events) + Subquery(total_transaction_events),
         plan_event_count=Cast(
             KeyTextTransform(
                 "events", "djstripe_customers__subscriptions__plan__product__metadata"
@@ -62,13 +40,15 @@ def warn_organization_throttle():
 
     # 80% to 100% of event quota
     queryset = queryset.filter(
-        event_count__gte=F("plan_event_count") * 0.80,
-        event_count__lte=F("plan_event_count"),
+        total_event_count__gte=F("plan_event_count") * 0.80,
+        total_event_count__lte=F("plan_event_count"),
     )
 
+    print("-----------")
+    print(queryset.query)
     for org in queryset:
         subscription = org.djstripe_customers.first().subscription
-        send_email_warn_quota.delay(subscription.pk, org.event_count)
+        send_email_warn_quota.delay(subscription.pk, org.total_event_count)
         warning, created = SubscriptionQuotaWarning.objects.get_or_create(
             subscription=subscription
         )

djstripe_ext/views.py

@@ -75,6 +75,7 @@ class SubscriptionViewSet(viewsets.ModelViewSet):
                     "eventCount": 0,
                     "transactionEventCount": 0,
                     "uptimeCheckEventCount": 0,
+                    "fileSizeMB": 0,
                 }
             )
         organization = subscription.customer.subscriber
@@ -83,7 +84,7 @@ class SubscriptionViewSet(viewsets.ModelViewSet):
         if data is None:
             org = Organization.objects.with_event_counts().get(pk=organization.pk)
             data = {
-                "eventCount": org.event_count,
+                "eventCount": org.issue_event_count,
                 "transactionEventCount": org.transaction_count,
                 "uptimeCheckEventCount": org.uptime_check_event_count,
                 "fileSizeMB": org.file_size,

organizations_ext/admin.py

@@ -1,6 +1,5 @@
 from django.conf import settings
 from django.contrib import admin
-from django.db.models import Count, OuterRef, Subquery, Sum
 from django.utils.html import format_html
 from organizations.base_admin import (
     BaseOrganizationAdmin,
@@ -8,8 +7,6 @@ from organizations.base_admin import (
     BaseOwnerInline,
 )
 
-from projects.models import Project
-
 from .models import Organization, OrganizationOwner, OrganizationUser
 
 ORGANIZATION_LIST_FILTER = (
@@ -47,7 +44,7 @@ class OrganizationAdmin(BaseOrganizationAdmin):
     show_full_result_count = False
 
     def issue_events(self, obj):
-        return obj.event_count
+        return obj.issue_event_count
 
     def customers(self, obj):
         return format_html(
@@ -66,7 +63,7 @@ class OrganizationAdmin(BaseOrganizationAdmin):
         return obj.uptime_check_event_count
 
     def total_events(self, obj):
-        return obj.event_count + obj.transaction_count + obj.uptime_check_event_count
+        return obj.total_event_count
 
     def get_queryset(self, request):
         queryset = self.model.objects.with_event_counts(False)

organizations_ext/models.py

@@ -184,7 +184,7 @@ class OrganizationManager(OrgManager):
             )
         ).values("total")
         return self.annotate(
-            event_count=Coalesce(Subquery(total_issue_events), 0),
+            issue_event_count=Coalesce(Subquery(total_issue_events), 0),
             transaction_count=Coalesce(Subquery(total_transaction_events), 0),
             uptime_check_event_count=Count(
                 "monitor__checks",
@@ -201,6 +201,10 @@ class OrganizationManager(OrgManager):
                 + Coalesce(total_dif_file_size, 0)
             )
             / 1000000,
+            total_event_count=F("issue_event_count")
+            + F("transaction_count")
+            + F("uptime_check_event_count")
+            + F("file_size"),
         )
 
 

organizations_ext/tasks.py

@@ -1,67 +1,33 @@
 from django.conf import settings
-from django.db.models import Count, F, Q, Subquery, OuterRef
-from django.db.models.functions import Coalesce
 from celery import shared_task
-from projects.models import Project
 from .models import Organization
 from .email import MetQuotaEmail, InvitationEmail
 
 
 def get_free_tier_organizations_with_event_count():
-    queryset = Organization.objects.filter(
+    return Organization.objects.with_event_counts().filter(
         djstripe_customers__subscriptions__plan__amount=0,
         djstripe_customers__subscriptions__status="active",
     )
 
-    projects = Project.objects.filter(organization=OuterRef("pk")).values(
-        "organization"
-    )
-    total_issue_events = projects.annotate(
-        total=Count(
-            "issue__event",
-            filter=Q(
-                issue__event__created__gte=OuterRef(
-                    "djstripe_customers__subscriptions__current_period_start"
-                )
-            ),
-        )
-    ).values("total")
-    total_transaction_events = projects.annotate(
-        total=Count(
-            "transactiongroup__transactionevent",
-            filter=Q(
-                transactiongroup__transactionevent__created__gte=OuterRef(
-                    "djstripe_customers__subscriptions__current_period_start"
-                )
-            ),
-        )
-    ).values("total")
-
-    return queryset.annotate(
-        uptime_check_event_count=Count("monitor__checks"),
-        event_count=Coalesce(Subquery(total_issue_events), 0)
-        + Coalesce(Subquery(total_transaction_events), 0)
-        + F("uptime_check_event_count"),
-    )
-
 
 @shared_task
 def set_organization_throttle():
-    """ Determine if organization should be throttled """
+    """Determine if organization should be throttled"""
     # Currently throttling only happens if billing is enabled and user has free plan.
     if settings.BILLING_ENABLED:
         events_max = settings.BILLING_FREE_TIER_EVENTS
         free_tier_organizations = get_free_tier_organizations_with_event_count()
 
         orgs_over_quota = free_tier_organizations.filter(
-            is_accepting_events=True, event_count__gt=events_max
+            is_accepting_events=True, total_event_count__gt=events_max
         ).select_related("owner__organization_user")
         for org in orgs_over_quota:
             send_email_met_quota.delay(org.pk)
         orgs_over_quota.update(is_accepting_events=False)
 
         free_tier_organizations.filter(
-            is_accepting_events=False, event_count__lte=events_max
+            is_accepting_events=False, total_event_count__lte=events_max
         ).update(is_accepting_events=True)
 
         # paid accounts should always be active at this time

organizations_ext/tests/test_throttling.py

@@ -28,6 +28,7 @@ class OrganizationThrottlingTestCase(TestCase):
                 livemode=False,
                 plan=plan,
                 status="active",
+                current_period_end="2000-01-31",
             )
             baker.make(
                 "events.Event", issue__project__organization=organization, _quantity=3
@@ -49,6 +50,9 @@ class OrganizationThrottlingTestCase(TestCase):
             subscription.current_period_start = timezone.make_aware(
                 timezone.datetime(2000, 2, 1)
             )
+            subscription.current_period_end = timezone.make_aware(
+                timezone.datetime(2000, 2, 28)
+            )
             subscription.save()
             set_organization_throttle()
             organization.refresh_from_db()
@@ -84,13 +88,16 @@ class OrganizationThrottlingTestCase(TestCase):
                 livemode=False,
                 plan=plan,
                 status="active",
+                current_period_end="2000-02-01",
             )
             baker.make("events.Event", issue__project=project, _quantity=3)
             baker.make(
-                "performance.TransactionEvent", group__project=project, _quantity=2,
+                "performance.TransactionEvent",
+                group__project=project,
+                _quantity=2,
             )
             free_org = get_free_tier_organizations_with_event_count().first()
-        self.assertEqual(free_org.event_count, 5)
+        self.assertEqual(free_org.total_event_count, 5)
 
     @override_settings(BILLING_FREE_TIER_EVENTS=1)
     def test_non_subscriber_throttling_performance(self):
@@ -112,5 +119,5 @@ class OrganizationThrottlingTestCase(TestCase):
             baker.make(
                 "events.Event", issue__project__organization=organization, _quantity=2
             )
-        with self.assertNumQueries(12):
+        with self.assertNumQueries(4):
             set_organization_throttle()