diff --git a/api/organisations/constants.py b/api/organisations/constants.py index 1fc68893ebce..26e9c93a2a9a 100644 --- a/api/organisations/constants.py +++ b/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)" ) diff --git a/api/organisations/migrations/0053_create_api_limit_access_block.py b/api/organisations/migrations/0053_create_api_limit_access_block.py new file mode 100644 index 000000000000..221f02dfb971 --- /dev/null +++ b/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')), + ], + ), + ] diff --git a/api/organisations/models.py b/api/organisations/models.py index c53825b0b229..1323b2c7162d 100644 --- a/api/organisations/models.py +++ b/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, diff --git a/api/organisations/tasks.py b/api/organisations/tasks.py index 279894ec3b2d..37701d68829f 100644 --- a/api/organisations/tasks.py +++ b/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,6 +32,7 @@ ALERT_EMAIL_MESSAGE, ALERT_EMAIL_SUBJECT, API_USAGE_ALERT_THRESHOLDS, + API_USAGE_GRACE_PERIOD, ) from .subscriptions.constants import SubscriptionCacheEntity @@ -36,7 +40,7 @@ @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) diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py index 437e9841ce55..3489417323ee 100644 --- a/api/tests/unit/organisations/test_unit_organisations_tasks.py +++ b/api/tests/unit/organisations/test_unit_organisations_tasks.py @@ -5,10 +5,13 @@ import pytest from django.core.mail.message import EmailMultiAlternatives from django.utils import timezone +from freezegun.api import FrozenDateTimeFactory from pytest_mock import MockerFixture from organisations.chargebee.metadata import ChargebeeObjMetadata +from organisations.constants import API_USAGE_GRACE_PERIOD from organisations.models import ( + APILimitAccessBlock, OranisationAPIUsageNotification, Organisation, OrganisationRole, @@ -16,6 +19,7 @@ UserOrganisation, ) from organisations.subscriptions.constants import ( + CHARGEBEE, FREE_PLAN_ID, MAX_API_CALLS_IN_FREE_PLAN, MAX_SEATS_IN_FREE_PLAN, @@ -26,8 +30,10 @@ ALERT_EMAIL_SUBJECT, finish_subscription_cancellation, handle_api_usage_notifications, + restrict_use_due_to_api_limit_grace_period_over, send_org_over_limit_alert, send_org_subscription_cancelled_alert, + unrestrict_after_api_limit_grace_period_is_stale, ) from users.models import FFAdminUser @@ -452,3 +458,194 @@ def test_handle_api_usage_notifications_above_100( == 1 ) assert OranisationAPIUsageNotification.objects.first() == api_usage_notification + + +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") +def test_restrict_use_due_to_api_limit_grace_period_over( + mocker: MockerFixture, + organisation: Organisation, + freezer: FrozenDateTimeFactory, +) -> None: + # Given + get_client_mock = mocker.patch("organisations.tasks.get_client") + client_mock = MagicMock() + get_client_mock.return_value = client_mock + client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + + now = timezone.now() + organisation2 = Organisation.objects.create(name="Org #2") + organisation3 = Organisation.objects.create(name="Org #3") + organisation4 = Organisation.objects.create(name="Org #4") + organisation5 = Organisation.objects.create(name="Org #5") + + organisation5.subscription.plan = "scale-up-v2" + organisation5.subscription.payment_method = CHARGEBEE + organisation5.subscription.subscription_id = "subscription-id" + organisation5.subscription.save() + + OranisationAPIUsageNotification.objects.create( + notified_at=now, + organisation=organisation, + percent_usage=100, + ) + OranisationAPIUsageNotification.objects.create( + notified_at=now, + organisation=organisation, + percent_usage=120, + ) + OranisationAPIUsageNotification.objects.create( + notified_at=now, + organisation=organisation2, + percent_usage=100, + ) + + # Should be ignored, since percent usage is less than 100. + OranisationAPIUsageNotification.objects.create( + notified_at=now, + organisation=organisation3, + percent_usage=90, + ) + + # Should be ignored, since not on a free plan. + OranisationAPIUsageNotification.objects.create( + notified_at=now, + organisation=organisation5, + percent_usage=120, + ) + + now = now + timedelta(days=API_USAGE_GRACE_PERIOD + 1) + freezer.move_to(now) + + # Should be ignored, since the notify period is too recent. + OranisationAPIUsageNotification.objects.create( + notified_at=now, + organisation=organisation3, + percent_usage=120, + ) + + # When + restrict_use_due_to_api_limit_grace_period_over() + + # Then + organisation.refresh_from_db() + organisation2.refresh_from_db() + organisation3.refresh_from_db() + organisation4.refresh_from_db() + organisation5.refresh_from_db() + + # Organisation without breaching 100 percent usage is ok. + assert organisation3.stop_serving_flags is False + assert organisation3.block_access_to_admin is False + assert getattr(organisation3, "api_limit_access_block", None) is None + + # Organisation which is still in the grace period is ok. + assert organisation4.stop_serving_flags is False + assert organisation4.block_access_to_admin is False + assert getattr(organisation4, "api_limit_access_block", None) is None + + # Organisation which is not on the free plan is ok. + assert organisation5.stop_serving_flags is False + assert organisation5.block_access_to_admin is False + assert getattr(organisation5, "api_limit_access_block", None) is None + + # Organisations that breached 100 are blocked. + assert organisation.stop_serving_flags is True + assert organisation.block_access_to_admin is True + assert organisation.api_limit_access_block + assert organisation2.stop_serving_flags is True + assert organisation2.block_access_to_admin is True + assert organisation2.api_limit_access_block + + # Organisations that change their subscription are unblocked. + organisation.subscription.plan = "scale-up-v2" + organisation.subscription.save() + organisation.refresh_from_db() + assert organisation.stop_serving_flags is False + assert organisation.block_access_to_admin is False + assert getattr(organisation, "api_limit_access_block", None) is None + + +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") +def test_unrestrict_after_api_limit_grace_period_is_stale( + organisation: Organisation, + freezer: FrozenDateTimeFactory, +) -> None: + # Given + now = timezone.now() + organisation2 = Organisation.objects.create(name="Org #2") + organisation3 = Organisation.objects.create(name="Org #3") + organisation4 = Organisation.objects.create(name="Org #4") + + organisation.stop_serving_flags = True + organisation.block_access_to_admin = True + organisation.save() + + organisation2.stop_serving_flags = True + organisation2.block_access_to_admin = True + organisation2.save() + + organisation3.stop_serving_flags = True + organisation3.block_access_to_admin = True + organisation3.save() + + organisation4.stop_serving_flags = True + organisation4.block_access_to_admin = True + organisation4.save() + + # Create access blocks for the first three, excluding the 4th. + APILimitAccessBlock.objects.create(organisation=organisation) + APILimitAccessBlock.objects.create(organisation=organisation2) + APILimitAccessBlock.objects.create(organisation=organisation3) + + OranisationAPIUsageNotification.objects.create( + notified_at=now, + organisation=organisation, + percent_usage=100, + ) + OranisationAPIUsageNotification.objects.create( + notified_at=now, + organisation=organisation, + percent_usage=120, + ) + OranisationAPIUsageNotification.objects.create( + notified_at=now, + organisation=organisation2, + percent_usage=100, + ) + + now = now + timedelta(days=32) + freezer.move_to(now) + + # Exclude the organisation since there's a recent notification. + OranisationAPIUsageNotification.objects.create( + notified_at=now, + organisation=organisation3, + percent_usage=120, + ) + + # When + unrestrict_after_api_limit_grace_period_is_stale() + + # Then + organisation.refresh_from_db() + organisation2.refresh_from_db() + organisation3.refresh_from_db() + organisation4.refresh_from_db() + + # Organisations with stale notifications revert to access. + assert organisation.stop_serving_flags is False + assert organisation.block_access_to_admin is False + assert organisation2.stop_serving_flags is False + assert organisation2.block_access_to_admin is False + assert getattr(organisation, "api_limit_access_block", None) is None + assert getattr(organisation2, "api_limit_access_block", None) is None + + # Organisations with recent API usage notifications are blocked. + assert organisation3.stop_serving_flags is True + assert organisation3.block_access_to_admin is True + assert organisation3.api_limit_access_block + + # Organisations without api limit access blocks stay blocked. + assert organisation4.stop_serving_flags is True + assert organisation4.block_access_to_admin is True + assert getattr(organisation4, "api_limit_access_block", None) is None