feat: Add API usage billing (#3729)

Co-authored-by: Matthew Elwell <[email protected]>

api/organisations/chargebee/__init__.py

@@ -1,4 +1,7 @@
 from .chargebee import (  # noqa
+    add_1000_api_calls,
+    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": False,
+                "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/0054_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', '0053_create_api_limit_access_block'),
+    ]
+
+    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

@@ -458,7 +458,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"
     )
@@ -490,3 +490,30 @@ 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):
+    """
+    Tracks API billing for when accounts go over their API usage
+    limits. This model is what allows subsequent billing runs
+    to not double bill an organisation for the same use.
+
+    Even though api_overage is charge per thousand API calls, this
+    class tracks the actual rounded count of API calls that are
+    billed for (i.e., 52000 for an account with 52233 api calls).
+    We're intentionally rounding down to the closest thousands.
+
+    The option to set immediate_invoice means whether or not the
+    API billing was processed immediately versus pushed onto the
+    subsequent subscription billing period.
+    """
+
+    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,16 +5,21 @@
 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.db.models import F, 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.chargebee import (
+    add_1000_api_calls_scale_up,
+    add_1000_api_calls_start_up,
+)
 from organisations.models import (
     APILimitAccessBlock,
-    OranisationAPIUsageNotification,
     Organisation,
+    OrganisationAPIBilling,
+    OrganisationAPIUsageNotification,
     OrganisationRole,
     Subscription,
 )
@@ -34,7 +39,13 @@
     API_USAGE_ALERT_THRESHOLDS,
     API_USAGE_GRACE_PERIOD,
 )
-from .subscriptions.constants import SubscriptionCacheEntity
+from .subscriptions.constants import (
+    SCALE_UP,
+    SCALE_UP_V2,
+    STARTUP,
+    STARTUP_V2,
+    SubscriptionCacheEntity,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -130,7 +141,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(),
@@ -162,7 +173,7 @@ def _handle_api_usage_notifications(organisation: Organisation) -> None:
     if matched_threshold is None:
         return
 
-    if OranisationAPIUsageNotification.objects.filter(
+    if OrganisationAPIUsageNotification.objects.filter(
         notified_at__gt=period_starts_at,
         percent_usage=matched_threshold,
     ).exists():
@@ -197,6 +208,85 @@ def handle_api_usage_notifications() -> None:
             )
 
 
+def charge_for_api_call_count_overages():
+    now = timezone.now()
+
+    # Get the period where we're interested in any new API usage
+    # notifications for the relevant billing period (ie, this month).
+    api_usage_notified_at = now - timedelta(days=30)
+
+    # Since we're only interested in monthly billed accounts, set a wide
+    # threshold to catch as many billing periods that could be roughly
+    # considered to be a "monthly" subscription, while still ruling out
+    # non-monthly subscriptions.
+    month_window_start = timedelta(days=25)
+    month_window_end = timedelta(days=35)
+
+    # Only apply charges to ongoing subscriptions that are close to
+    # being charged due to being at the end of the billing term.
+    closing_billing_term = now + timedelta(hours=1)
+
+    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",
+    ):
+        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,
+        )
+
+
 def restrict_use_due_to_api_limit_grace_period_over() -> None:
     """
     Restrict API use once a grace period has ended.
@@ -208,7 +298,7 @@ def restrict_use_due_to_api_limit_grace_period_over() -> None:
     grace_period = timezone.now() - timedelta(days=API_USAGE_GRACE_PERIOD)
     month_start = timezone.now() - timedelta(30)
     queryset = (
-        OranisationAPIUsageNotification.objects.filter(
+        OrganisationAPIUsageNotification.objects.filter(
             notified_at__gt=month_start,
             notified_at__lt=grace_period,
             percent_usage__gte=100,
@@ -270,7 +360,7 @@ def unrestrict_after_api_limit_grace_period_is_stale() -> None:
 
     month_start = timezone.now() - timedelta(30)
     still_restricted_organisation_notifications = (
-        OranisationAPIUsageNotification.objects.filter(
+        OrganisationAPIUsageNotification.objects.filter(
             notified_at__gt=month_start,
             percent_usage__gte=100,
         )
@@ -296,13 +386,28 @@ def unrestrict_after_api_limit_grace_period_is_stale() -> None:
         organisation.api_limit_access_block.delete()
 
 
-if settings.ENABLE_API_USAGE_ALERTING:
+def register_recurring_tasks() -> None:
+    """
+    Helper function to get codecov coverage.
+    """
+    assert settings.ENABLE_API_USAGE_ALERTING
+
     register_recurring_task(
         run_every=timedelta(hours=12),
     )(handle_api_usage_notifications)
+
+    register_recurring_task(
+        run_every=timedelta(minutes=30),
+    )(charge_for_api_call_count_overages)
+
     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)
+
+
+if settings.ENABLE_API_USAGE_ALERTING:
+    register_recurring_tasks()  # pragma: no cover