feat: Block access after seven days notice of API overage (#3714)

api/organisations/constants.py

@@ -1,4 +1,5 @@
 API_USAGE_ALERT_THRESHOLDS = [75, 90, 100, 120]
+API_USAGE_GRACE_PERIOD = 7
 ALERT_EMAIL_MESSAGE = (
     "Organisation %s has used %d seats which is over their plan limit of %d (plan: %s)"
 )

api/organisations/migrations/0053_create_api_limit_access_block.py

@@ -0,0 +1,23 @@
+# Generated by Django 3.2.25 on 2024-04-03 14:11
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('organisations', '0052_create_hubspot_organisation'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='APILimitAccessBlock',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('created_at', models.DateTimeField(auto_now_add=True, null=True)),
+                ('updated_at', models.DateTimeField(auto_now=True, null=True)),
+                ('organisation', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, related_name='api_limit_access_block', to='organisations.organisation')),
+            ],
+        ),
+    ]

api/organisations/models.py

@@ -254,6 +254,16 @@ def can_auto_upgrade_seats(self) -> bool:
     def is_free_plan(self) -> bool:
         return self.plan == FREE_PLAN_ID
 
+    @hook(AFTER_SAVE, when="plan", has_changed=True)
+    def update_api_limit_access_block(self):
+        if not getattr(self.organisation, "api_limit_access_block", None):
+            return
+
+        self.organisation.api_limit_access_block.delete()
+        self.organisation.stop_serving_flags = False
+        self.organisation.block_access_to_admin = False
+        self.organisation.save()
+
     @hook(AFTER_SAVE, when="plan", has_changed=True)
     def update_hubspot_active_subscription(self):
         if not settings.ENABLE_HUBSPOT_LEAD_TRACKING:
@@ -462,6 +472,15 @@ class OranisationAPIUsageNotification(models.Model):
     updated_at = models.DateTimeField(null=True, auto_now=True)
 
 
+class APILimitAccessBlock(models.Model):
+    organisation = models.OneToOneField(
+        Organisation, on_delete=models.CASCADE, related_name="api_limit_access_block"
+    )
+
+    created_at = models.DateTimeField(null=True, auto_now_add=True)
+    updated_at = models.DateTimeField(null=True, auto_now=True)
+
+
 class HubspotOrganisation(models.Model):
     organisation = models.OneToOneField(
         Organisation,

api/organisations/tasks.py

@@ -5,17 +5,20 @@
 from dateutil.relativedelta import relativedelta
 from django.conf import settings
 from django.core.mail import send_mail
+from django.db.models import Max
 from django.template.loader import render_to_string
 from django.utils import timezone
 
 from integrations.flagsmith.client import get_client
 from organisations import subscription_info_cache
 from organisations.models import (
+    APILimitAccessBlock,
     OranisationAPIUsageNotification,
     Organisation,
     OrganisationRole,
     Subscription,
 )
+from organisations.subscriptions.constants import FREE_PLAN_ID
 from organisations.subscriptions.subscription_service import (
     get_subscription_metadata,
 )
@@ -29,14 +32,15 @@
     ALERT_EMAIL_MESSAGE,
     ALERT_EMAIL_SUBJECT,
     API_USAGE_ALERT_THRESHOLDS,
+    API_USAGE_GRACE_PERIOD,
 )
 from .subscriptions.constants import SubscriptionCacheEntity
 
 logger = logging.getLogger(__name__)
 
 
 @register_task_handler()
-def send_org_over_limit_alert(organisation_id):
+def send_org_over_limit_alert(organisation_id) -> None:
     organisation = Organisation.objects.get(id=organisation_id)
 
     subscription_metadata = get_subscription_metadata(organisation)
@@ -56,7 +60,7 @@ def send_org_over_limit_alert(organisation_id):
 def send_org_subscription_cancelled_alert(
     organisation_name: str,
     formatted_cancellation_date: str,
-):
+) -> None:
     FFAdminUser.send_alert_to_admin_users(
         subject=f"Organisation {organisation_name} has cancelled their subscription",
         message=f"Organisation {organisation_name} has cancelled their subscription on {formatted_cancellation_date}",
@@ -69,7 +73,7 @@ def update_organisation_subscription_information_influx_cache():
 
 
 @register_task_handler()
-def update_organisation_subscription_information_cache():
+def update_organisation_subscription_information_cache() -> None:
     subscription_info_cache.update_caches(
         (SubscriptionCacheEntity.CHARGEBEE, SubscriptionCacheEntity.INFLUX)
     )
@@ -78,7 +82,7 @@ def update_organisation_subscription_information_cache():
 @register_recurring_task(
     run_every=timedelta(hours=12),
 )
-def finish_subscription_cancellation():
+def finish_subscription_cancellation() -> None:
     now = timezone.now()
     previously = now + timedelta(hours=-24)
     for subscription in Subscription.objects.filter(
@@ -133,7 +137,7 @@ def send_admin_api_usage_notification(
     )
 
 
-def _handle_api_usage_notifications(organisation: Organisation):
+def _handle_api_usage_notifications(organisation: Organisation) -> None:
     subscription_cache = organisation.subscription_information_cache
     billing_starts_at = subscription_cache.current_billing_term_starts_at
     now = timezone.now()
@@ -154,6 +158,10 @@ def _handle_api_usage_notifications(organisation: Organisation):
 
         matched_threshold = threshold
 
+    # Didn't match even the lowest threshold, so no notification.
+    if matched_threshold is None:
+        return
+
     if OranisationAPIUsageNotification.objects.filter(
         notified_at__gt=period_starts_at,
         percent_usage=matched_threshold,
@@ -164,7 +172,7 @@ def _handle_api_usage_notifications(organisation: Organisation):
     send_admin_api_usage_notification(organisation, matched_threshold)
 
 
-def handle_api_usage_notifications():
+def handle_api_usage_notifications() -> None:
     flagsmith_client = get_client("local", local_eval=True)
 
     for organisation in Organisation.objects.filter(
@@ -189,7 +197,112 @@ def handle_api_usage_notifications():
             )
 
 
+def restrict_use_due_to_api_limit_grace_period_over() -> None:
+    """
+    Restrict API use once a grace period has ended.
+
+    Since free plans don't have predefined subscription periods, we
+    use a rolling thirty day period to filter them.
+    """
+
+    grace_period = timezone.now() - timedelta(days=API_USAGE_GRACE_PERIOD)
+    month_start = timezone.now() - timedelta(30)
+    queryset = (
+        OranisationAPIUsageNotification.objects.filter(
+            notified_at__gt=month_start,
+            notified_at__lt=grace_period,
+            percent_usage__gte=100,
+        )
+        .values("organisation")
+        .annotate(max_value=Max("percent_usage"))
+    )
+
+    organisation_ids = []
+    for result in queryset:
+        organisation_ids.append(result["organisation"])
+    organisations = Organisation.objects.filter(
+        id__in=organisation_ids,
+        subscription__plan=FREE_PLAN_ID,
+        api_limit_access_block__isnull=True,
+    ).exclude(
+        stop_serving_flags=True,
+        block_access_to_admin=True,
+    )
+
+    update_organisations = []
+    api_limit_access_blocks = []
+    flagsmith_client = get_client("local", local_eval=True)
+
+    for organisation in organisations:
+        flags = flagsmith_client.get_identity_flags(
+            f"org.{organisation.id}.{organisation.name}",
+            traits={"organisation_id": organisation.id},
+        )
+
+        stop_serving = flags.is_feature_enabled("api_limiting_stop_serving_flags")
+        block_access = flags.is_feature_enabled("api_limiting_block_access_to_admin")
+
+        if not stop_serving and not block_access:
+            continue
+
+        organisation.stop_serving_flags = stop_serving
+        organisation.block_access_to_admin = block_access
+
+        api_limit_access_blocks.append(APILimitAccessBlock(organisation=organisation))
+        update_organisations.append(organisation)
+
+    APILimitAccessBlock.objects.bulk_create(api_limit_access_blocks)
+
+    Organisation.objects.bulk_update(
+        update_organisations, ["stop_serving_flags", "block_access_to_admin"]
+    )
+
+
+def unrestrict_after_api_limit_grace_period_is_stale() -> None:
+    """
+    This task handles accounts that have breached the API limit
+    and have become restricted by setting the stop_serving_flags
+    and block_access_to_admin to True. This task looks to find
+    which accounts have started following the API limits in the
+    latest rolling month and re-enables them if they no longer
+    have recent API usage notifications.
+    """
+
+    month_start = timezone.now() - timedelta(30)
+    still_restricted_organisation_notifications = (
+        OranisationAPIUsageNotification.objects.filter(
+            notified_at__gt=month_start,
+            percent_usage__gte=100,
+        )
+        .values("organisation")
+        .annotate(max_value=Max("percent_usage"))
+    )
+    still_restricted_organisation_ids = {
+        q["organisation"] for q in still_restricted_organisation_notifications
+    }
+    organisation_ids = set(
+        Organisation.objects.filter(
+            api_limit_access_block__isnull=False,
+        ).values_list("id", flat=True)
+    )
+
+    matching_organisations = Organisation.objects.filter(
+        id__in=(organisation_ids - still_restricted_organisation_ids),
+    )
+
+    matching_organisations.update(stop_serving_flags=False, block_access_to_admin=False)
+
+    for organisation in matching_organisations:
+        organisation.api_limit_access_block.delete()
+
+
 if settings.ENABLE_API_USAGE_ALERTING:
     register_recurring_task(
         run_every=timedelta(hours=12),
     )(handle_api_usage_notifications)
+    register_recurring_task(
+        run_every=timedelta(hours=12),
+    )(restrict_use_due_to_api_limit_grace_period_over)
+    register_recurring_task(
+        run_every=timedelta(hours=12),
+    )(unrestrict_after_api_limit_grace_period_is_stale)