feat: Add API usage billing

api/organisations/chargebee/__init__.py

@@ -1,4 +1,6 @@
 from .chargebee import (  # noqa
+    add_1000_api_calls_scale_up,
+    add_1000_api_calls_start_up,
     add_single_seat,
     extract_subscription_metadata,
     get_customer_id_from_subscription_id,

api/organisations/chargebee/chargebee.py

@@ -11,11 +11,17 @@
 from ..subscriptions.constants import CHARGEBEE
 from ..subscriptions.exceptions import (
     CannotCancelChargebeeSubscription,
+    UpgradeAPIUsageError,
+    UpgradeAPIUsagePaymentFailure,
     UpgradeSeatsError,
     UpgradeSeatsPaymentFailure,
 )
 from .cache import ChargebeeCache
-from .constants import ADDITIONAL_SEAT_ADDON_ID
+from .constants import (
+    ADDITIONAL_API_SCALE_UP_ADDON_ID,
+    ADDITIONAL_API_START_UP_ADDON_ID,
+    ADDITIONAL_SEAT_ADDON_ID,
+)
 from .metadata import ChargebeeObjMetadata
 
 chargebee.configure(settings.CHARGEBEE_API_KEY, settings.CHARGEBEE_SITE)
@@ -203,6 +209,54 @@ def add_single_seat(subscription_id: str):
         raise UpgradeSeatsError(msg) from e
 
 
+def add_1000_api_calls_start_up(
+    subscription_id: str, count: int = 1, invoice_immediately: bool = False
+) -> None:
+    add_1000_api_calls(ADDITIONAL_API_START_UP_ADDON_ID, subscription_id, count)
+
+
+def add_1000_api_calls_scale_up(
+    subscription_id: str, count: int = 1, invoice_immediately: bool = False
+) -> None:
+    add_1000_api_calls(ADDITIONAL_API_SCALE_UP_ADDON_ID, subscription_id, count)
+
+
+def add_1000_api_calls(
+    addon_id: str,
+    subscription_id: str,
+    count: int = 1,
+    invoice_immediately: bool = False,
+) -> None:
+    if not count:
+        return
+    try:
+        chargebee.Subscription.update(
+            subscription_id,
+            {
+                "addons": [{"id": addon_id, "quantity": count}],
+                "prorate": True,
+                "invoice_immediately": invoice_immediately,
+            },
+        )
+
+    except ChargebeeAPIError as e:
+        api_error_code = e.json_obj["api_error_code"]
+        if api_error_code in CHARGEBEE_PAYMENT_ERROR_CODES:
+            logger.warning(
+                f"Payment declined ({api_error_code}) during additional "
+                f"api calls upgrade to a CB subscription for subscription_id "
+                f"{subscription_id}"
+            )
+            raise UpgradeAPIUsagePaymentFailure() from e
+
+        msg = (
+            "Failed to add additional API calls to CB subscription for subscription id: %s"
+            % subscription_id
+        )
+        logger.error(msg)
+        raise UpgradeAPIUsageError(msg) from e
+
+
 def _convert_chargebee_subscription_to_dictionary(
     chargebee_subscription: chargebee.Subscription,
 ) -> dict:

api/organisations/chargebee/constants.py

@@ -1 +1,4 @@
 ADDITIONAL_SEAT_ADDON_ID = "additional-team-members-scale-up-v2-monthly"
+
+ADDITIONAL_API_START_UP_ADDON_ID = "additional-api-start-up-monthly"
+ADDITIONAL_API_SCALE_UP_ADDON_ID = "additional-api-scale-up-monthly"

api/organisations/migrations/0053_create_api_billing.py

@@ -0,0 +1,30 @@
+# Generated by Django 3.2.25 on 2024-04-08 15:07
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('organisations', '0052_create_hubspot_organisation'),
+    ]
+
+    operations = [
+        migrations.RenameModel(
+            old_name='OranisationAPIUsageNotification',
+            new_name='OrganisationAPIUsageNotification',
+        ),
+        migrations.CreateModel(
+            name='OrganisationAPIBilling',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('api_overage', models.IntegerField()),
+                ('immediate_invoice', models.BooleanField(default=False)),
+                ('billed_at', models.DateTimeField()),
+                ('created_at', models.DateTimeField(auto_now_add=True, null=True)),
+                ('updated_at', models.DateTimeField(auto_now=True, null=True)),
+                ('organisation', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='api_billing', to='organisations.organisation')),
+            ],
+        ),
+    ]

api/organisations/models.py

@@ -455,7 +455,7 @@ def erase_api_notifications(self):
         self.organisation.api_usage_notifications.all().delete()
 
 
-class OranisationAPIUsageNotification(models.Model):
+class OrganisationAPIUsageNotification(models.Model):
     organisation = models.ForeignKey(
         Organisation, on_delete=models.CASCADE, related_name="api_usage_notifications"
     )
@@ -478,3 +478,15 @@ class HubspotOrganisation(models.Model):
     hubspot_id = models.CharField(max_length=100)
     created_at = models.DateTimeField(auto_now_add=True)
     updated_at = models.DateTimeField(auto_now=True)
+
+
+class OrganisationAPIBilling(models.Model):
+    organisation = models.ForeignKey(
+        Organisation, on_delete=models.CASCADE, related_name="api_billing"
+    )
+    api_overage = models.IntegerField(null=False)
+    immediate_invoice = models.BooleanField(null=False, default=False)
+    billed_at = models.DateTimeField(null=False)
+
+    created_at = models.DateTimeField(null=True, auto_now_add=True)
+    updated_at = models.DateTimeField(null=True, auto_now=True)

api/organisations/subscriptions/constants.py

@@ -40,6 +40,13 @@
 )
 FREE_PLAN_ID = "free"
 TRIAL_SUBSCRIPTION_ID = "trial"
+SCALE_UP = "scale-up"
+SCALE_UP_12_MONTHS_V2 = "scale-up-12-months-v2"
+SCALE_UP_QUARTERLY_V2_SEMIANNUAL = "scale-up-quarterly-v2-semiannual"
+SCALE_UP_V2 = "scale-up-v2"
+STARTUP = "startup"
+STARTUP_ANNUAL_V2 = "startup-annual-v2"
+STARTUP_V2 = "startup-v2"
 
 
 class SubscriptionCacheEntity(Enum):

api/organisations/subscriptions/exceptions.py

@@ -22,6 +22,18 @@ class UpgradeSeatsPaymentFailure(APIException):
     )
 
 
+class UpgradeAPIUsageError(APIException):
+    default_detail = "Failed to upgrade API use in Chargebee"
+
+
+class UpgradeAPIUsagePaymentFailure(APIException):
+    status_code = 400
+    default_detail = (
+        "API usage upgrade has failed due to a payment issue. "
+        "If this persists, contact the organisation admin."
+    )
+
+
 class SubscriptionDoesNotSupportSeatUpgrade(APIException):
     status_code = 400
     default_detail = "Please Upgrade your plan to add additional seats/users"

api/organisations/tasks.py

@@ -5,14 +5,20 @@
 from dateutil.relativedelta import relativedelta
 from django.conf import settings
 from django.core.mail import send_mail
+from django.db.models import F
 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.chargebee import (
+    add_1000_api_calls_scale_up,
+    add_1000_api_calls_start_up,
+)
 from organisations.models import (
-    OranisationAPIUsageNotification,
     Organisation,
+    OrganisationAPIBilling,
+    OrganisationAPIUsageNotification,
     OrganisationRole,
     Subscription,
 )
@@ -30,7 +36,13 @@
     ALERT_EMAIL_SUBJECT,
     API_USAGE_ALERT_THRESHOLDS,
 )
-from .subscriptions.constants import SubscriptionCacheEntity
+from .subscriptions.constants import (
+    SCALE_UP,
+    SCALE_UP_V2,
+    STARTUP,
+    STARTUP_V2,
+    SubscriptionCacheEntity,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -126,7 +138,7 @@ def send_admin_api_usage_notification(
         fail_silently=True,
     )
 
-    OranisationAPIUsageNotification.objects.create(
+    OrganisationAPIUsageNotification.objects.create(
         organisation=organisation,
         percent_usage=matched_threshold,
         notified_at=timezone.now(),
@@ -154,7 +166,7 @@ def _handle_api_usage_notifications(organisation: Organisation):
 
         matched_threshold = threshold
 
-    if OranisationAPIUsageNotification.objects.filter(
+    if OrganisationAPIUsageNotification.objects.filter(
         notified_at__gt=period_starts_at,
         percent_usage=matched_threshold,
     ).exists():
@@ -189,7 +201,76 @@ def handle_api_usage_notifications():
             )
 
 
+def charge_for_api_call_count_overages():
+    now = timezone.now()
+    api_usage_notified_at = now - timedelta(days=30)
+    month_window_start = timedelta(days=25)
+    month_window_end = timedelta(days=35)
+    closing_billing_term = now + timedelta(hours=12)
+    organisation_ids = set(
+        OrganisationAPIUsageNotification.objects.filter(
+            notified_at__gte=api_usage_notified_at,
+            percent_usage__gte=100,
+        )
+        .exclude(
+            organisation__api_billing__billed_at__gt=api_usage_notified_at,
+        )
+        .values_list("organisation_id", flat=True)
+    )
+
+    for organisation in Organisation.objects.filter(
+        id__in=organisation_ids,
+        subscription_information_cache__current_billing_term_ends_at__lte=closing_billing_term,
+        subscription_information_cache__current_billing_term_ends_at__gte=now,
+        subscription_information_cache__current_billing_term_starts_at__lte=F(
+            "subscription_information_cache__current_billing_term_ends_at"
+        )
+        - month_window_start,
+        subscription_information_cache__current_billing_term_starts_at__gte=F(
+            "subscription_information_cache__current_billing_term_ends_at"
+        )
+        - month_window_end,
+    ).select_related(
+        "subscription_information_cache",
+    ):
+        subscription_cache = organisation.subscription_information_cache
+        api_usage = get_current_api_usage(organisation.id, "30d")
+        api_usage_ratio = api_usage / subscription_cache.allowed_30d_api_calls
+
+        if api_usage_ratio < 1.0:
+            logger.warning("API Usage does not match API Notification")
+            continue
+
+        api_overage = api_usage - subscription_cache.allowed_30d_api_calls
+
+        if organisation.subscription.plan in {SCALE_UP, SCALE_UP_V2}:
+            add_1000_api_calls_scale_up(
+                organisation.subscription.subscription_id, api_overage // 1000
+            )
+        elif organisation.subscription.plan in {STARTUP, STARTUP_V2}:
+            add_1000_api_calls_start_up(
+                organisation.subscription.subscription_id, api_overage // 1000
+            )
+        else:
+            logger.error(
+                f"Unable to bill for API overages for plan `{organisation.subscription.plan}`"
+            )
+            continue
+
+        # Save a copy of what was just billed in order to avoid
+        # double billing on a subsequent task run.
+        OrganisationAPIBilling.objects.create(
+            organisation=organisation,
+            api_overage=(1000 * (api_overage // 1000)),
+            immediate_invoice=False,
+            billed_at=now,
+        )
+
+
 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),
+    )(charge_for_api_call_count_overages)

api/organisations/views.py

@@ -24,8 +24,8 @@
 from organisations.chargebee import webhook_event_types, webhook_handlers
 from organisations.exceptions import OrganisationHasNoPaidSubscription
 from organisations.models import (
-    OranisationAPIUsageNotification,
     Organisation,
+    OrganisationAPIUsageNotification,
     OrganisationRole,
     OrganisationWebhook,
 )
@@ -319,15 +319,15 @@ class OrganisationAPIUsageNotificationView(ListAPIView):
     def get_queryset(self):
         organisation = Organisation.objects.get(id=self.kwargs["organisation_pk"])
         if not hasattr(organisation, "subscription_information_cache"):
-            return OranisationAPIUsageNotification.objects.none()
+            return OrganisationAPIUsageNotification.objects.none()
         subscription_cache = organisation.subscription_information_cache
         billing_starts_at = subscription_cache.current_billing_term_starts_at
         now = timezone.now()
 
         month_delta = relativedelta(now, billing_starts_at).months
         period_starts_at = relativedelta(months=month_delta) + billing_starts_at
 
-        queryset = OranisationAPIUsageNotification.objects.filter(
+        queryset = OrganisationAPIUsageNotification.objects.filter(
             organisation_id=organisation.id,
             notified_at__gt=period_starts_at,
         )

api/tests/unit/organisations/test_unit_organisations_models.py

@@ -10,8 +10,8 @@
 from environments.models import Environment
 from organisations.chargebee.metadata import ChargebeeObjMetadata
 from organisations.models import (
-    OranisationAPIUsageNotification,
     Organisation,
+    OrganisationAPIUsageNotification,
     OrganisationSubscriptionInformationCache,
     Subscription,
 )
@@ -551,15 +551,15 @@ def test_reset_of_api_notifications(organisation: Organisation) -> None:
     )
 
     # Create a notification which should be deleted shortly.
-    OranisationAPIUsageNotification.objects.create(
+    OrganisationAPIUsageNotification.objects.create(
         organisation=organisation,
         percent_usage=90,
         notified_at=now,
     )
 
     # Keep a notification which should not be deleted.
     organisation2 = Organisation.objects.create(name="Test org2")
-    oapiun = OranisationAPIUsageNotification.objects.create(
+    oapiun = OrganisationAPIUsageNotification.objects.create(
         organisation=organisation2,
         percent_usage=90,
         notified_at=now,
@@ -570,5 +570,5 @@ def test_reset_of_api_notifications(organisation: Organisation) -> None:
     osic.save()
 
     # Then
-    assert OranisationAPIUsageNotification.objects.count() == 1
-    assert OranisationAPIUsageNotification.objects.first() == oapiun
+    assert OrganisationAPIUsageNotification.objects.count() == 1
+    assert OrganisationAPIUsageNotification.objects.first() == oapiun