GlitchTip/apps/organizations_ext/tasks.py

from celery import shared_task
from django.core.cache import cache

from .email import InvitationEmail, ThrottleNoticeEmail
from .models import Organization


@shared_task
def check_organization_throttle(organization_id: int):
    if not cache.add(f"org-throttle-{organization_id}", True):
        return  # Recent check already performed

    org = Organization.objects.with_event_counts().get(id=organization_id)
    _check_and_update_throttle(org)


@shared_task
def check_all_organizations_throttle():
    for org in Organization.objects.with_event_counts().iterator():
        _check_and_update_throttle(org)


def _check_and_update_throttle(org: Organization):
    from djstripe.models import Product

    plan_events: int | None = (
        Product.objects.filter(
            plan__subscriptions__customer__subscriber=org,
            plan__subscriptions__status="active",
        )
        .values_list("metadata__events", flat=True)
        .first()
    )
    if plan_events:
        plan_events = int(plan_events)
    org_throttle = 0
    if plan_events is None or org.total_event_count > plan_events * 2:
        org_throttle = 100
    elif org.total_event_count > plan_events * 1.5:
        org_throttle = 50
    elif org.total_event_count > plan_events:
        org_throttle = 10

    if org.event_throttle_rate != org_throttle:
        old_throttle = org.event_throttle_rate
        org.event_throttle_rate = org_throttle
        org.save(update_fields=["event_throttle_rate"])
        if org_throttle > old_throttle:
            send_throttle_email.delay(org.id)


@shared_task
def send_throttle_email(organization_id: int):
    ThrottleNoticeEmail(pk=organization_id).send_email()


@shared_task
def send_email_invite(org_user_id: int, token: str):
    InvitationEmail(pk=org_user_id, token=token).send_email()