fix: optimise influxdb calls from handle_api_usage_notifications task (#6458)

Author: matthewelwell

api/organisations/task_helpers.py

@@ -99,6 +99,7 @@ def _send_api_usage_notification(
 
 def handle_api_usage_notification_for_organisation(organisation: Organisation) -> None:
     now = timezone.now()
+    subscription_cache = organisation.subscription_information_cache
 
     if (
         organisation.subscription.is_free_plan
@@ -108,15 +109,7 @@ def handle_api_usage_notification_for_organisation(organisation: Organisation) -
         # Default to a rolling month for free accounts or canceled subscriptions.
         days = 30
         period_starts_at = now - timedelta(days)
-    elif not organisation.has_subscription_information_cache():
-        # Since the calling code is a list of many organisations
-        # log the error and return without raising an exception.
-        logger.error(
-            f"Paid organisation {organisation.id} is missing subscription information cache"
-        )
-        return
     else:
-        subscription_cache = organisation.subscription_information_cache
         billing_starts_at = subscription_cache.current_billing_term_starts_at
 
         if billing_starts_at is None:

api/organisations/tasks.py

@@ -33,6 +33,7 @@
 from .constants import (
     ALERT_EMAIL_MESSAGE,
     ALERT_EMAIL_SUBJECT,
+    API_USAGE_ALERT_THRESHOLDS,
     API_USAGE_GRACE_PERIOD,
 )
 from .subscriptions.constants import (
@@ -121,9 +122,20 @@ def finish_subscription_cancellation() -> None:
 def handle_api_usage_notifications() -> None:
     flagsmith_client = get_client("local", local_eval=True)
 
-    for organisation in Organisation.objects.all().select_related(
-        "subscription", "subscription_information_cache"
-    ):
+    threshold_percentage = min(API_USAGE_ALERT_THRESHOLDS) / 100
+    for organisation in Organisation.objects.filter(
+        # Not having a subscription information cache implies that the organisation has
+        # no usage so no notification is needed.
+        subscription_information_cache__isnull=False,
+        # Skip organisations whose usage in the last 30d is below the threshold of the
+        # minimum notification level for their allowance. If their usage is below the
+        # threshold in the last 30d, it must be below it for the current billing period
+        # (which by definition is <30d).
+        subscription_information_cache__api_calls_30d__gte=F(
+            "subscription_information_cache__allowed_30d_api_calls"
+        )
+        * threshold_percentage,
+    ).select_related("subscription", "subscription_information_cache"):
         feature_enabled = flagsmith_client.get_identity_flags(
             organisation.flagsmith_identifier,
             traits=organisation.flagsmith_on_flagsmith_api_traits,

api/tests/unit/organisations/test_unit_organisations_tasks.py

@@ -319,12 +319,22 @@ def test_handle_api_usage_notification_for_organisation_when_cancellation_date_i
     mocker: MockerFixture,
 ) -> None:
     # Given
+    usage = 25
+    allowance = 1_000_000
+
     organisation.subscription.plan = SCALE_UP
     organisation.subscription.subscription_id = "fancy_id"
     organisation.subscription.cancellation_date = timezone.now()
     organisation.subscription.save()
     mock_api_usage = mocker.patch("organisations.task_helpers.get_current_api_usage")
-    mock_api_usage.return_value = 25
+    mock_api_usage.return_value = usage
+
+    OrganisationSubscriptionInformationCache.objects.create(
+        organisation=organisation,
+        allowed_30d_api_calls=allowance,
+        api_calls_30d=usage,
+    )
+
     from organisations.task_helpers import logger
 
     logger.addHandler(inspecting_handler)
@@ -353,10 +363,11 @@ def test_handle_api_usage_notification_for_organisation_when_billing_starts_at_i
         organisation=organisation,
         allowed_30d_api_calls=1_000_000,
         current_billing_term_starts_at=billing_term_starts_at,
+        api_calls_30d=1_900_000,
     )
 
     mock_api_usage = mocker.patch("organisations.task_helpers.get_current_api_usage")
-    mock_api_usage.return_value = 25
+    mock_api_usage.return_value = 2_000_000
 
     organisation.refresh_from_db()
 
@@ -379,12 +390,11 @@ def test_handle_api_usage_notifications_when_feature_flag_is_off(
     now = timezone.now()
     OrganisationSubscriptionInformationCache.objects.create(
         organisation=organisation,
-        allowed_seats=10,
-        allowed_projects=3,
         allowed_30d_api_calls=100,
         chargebee_email="test@example.com",
         current_billing_term_starts_at=now - timedelta(days=45),
         current_billing_term_ends_at=now + timedelta(days=320),
+        api_calls_30d=110,
     )
     mock_api_usage = mocker.patch(
         "organisations.tasks.get_current_api_usage",
@@ -431,12 +441,10 @@ def test_handle_api_usage_notifications_below_100(
     organisation.subscription.save()
     OrganisationSubscriptionInformationCache.objects.create(
         organisation=organisation,
-        allowed_seats=10,
-        allowed_projects=3,
         allowed_30d_api_calls=100,
-        chargebee_email="test@example.com",
         current_billing_term_starts_at=now - timedelta(days=45),
         current_billing_term_ends_at=now + timedelta(days=320),
+        api_calls_30d=90,
     )
     mock_api_usage = mocker.patch(
         "organisations.task_helpers.get_current_api_usage",
@@ -465,7 +473,7 @@ def test_handle_api_usage_notifications_below_100(
     handle_api_usage_notifications()
 
     # Then
-    assert len(mock_api_usage.call_args_list) == 2
+    assert len(mock_api_usage.call_args_list) == 1
 
     # We only care about the call for the main organisation,
     # not the call for 'another_organisation'
@@ -551,6 +559,7 @@ def test_handle_api_usage_notifications_below_api_usage_alert_thresholds(
         chargebee_email="test@example.com",
         current_billing_term_starts_at=now - timedelta(days=45),
         current_billing_term_ends_at=now + timedelta(days=320),
+        api_calls_30d=70,
     )
     mock_api_usage = mocker.patch(
         "organisations.task_helpers.get_current_api_usage",
@@ -571,8 +580,6 @@ def test_handle_api_usage_notifications_below_api_usage_alert_thresholds(
     handle_api_usage_notifications()
 
     # Then
-    mock_api_usage.assert_called_once_with(organisation.id, now - timedelta(days=14))
-
     assert len(mailoutbox) == 0
 
     assert (
@@ -590,23 +597,25 @@ def test_handle_api_usage_notifications_above_100(
     mailoutbox: list[EmailMultiAlternatives],
 ) -> None:
     # Given
+    usage = 105
+    allowance = 100
+
     now = timezone.now()
     organisation.subscription.plan = SCALE_UP
     organisation.subscription.subscription_id = "fancy_id"
     organisation.subscription.save()
     OrganisationSubscriptionInformationCache.objects.create(
         organisation=organisation,
-        allowed_seats=10,
-        allowed_projects=3,
-        allowed_30d_api_calls=100,
+        allowed_30d_api_calls=allowance,
         chargebee_email="test@example.com",
         current_billing_term_starts_at=now - timedelta(days=45),
         current_billing_term_ends_at=now + timedelta(days=320),
+        api_calls_30d=usage,
     )
     mock_api_usage = mocker.patch(
         "organisations.task_helpers.get_current_api_usage",
     )
-    mock_api_usage.return_value = 105
+    mock_api_usage.return_value = usage
 
     get_client_mock = mocker.patch("organisations.tasks.get_client")
     client_mock = MagicMock()
@@ -693,12 +702,11 @@ def test_handle_api_usage_notifications_with_error(
     organisation.subscription.save()
     OrganisationSubscriptionInformationCache.objects.create(
         organisation=organisation,
-        allowed_seats=10,
-        allowed_projects=3,
         allowed_30d_api_calls=100,
         chargebee_email="test@example.com",
         current_billing_term_starts_at=now - timedelta(days=45),
         current_billing_term_ends_at=now + timedelta(days=320),
+        api_calls_30d=100,
     )
 
     get_client_mock = mocker.patch("organisations.tasks.get_client")
@@ -738,15 +746,24 @@ def test_handle_api_usage_notifications_for_free_accounts(
     mailoutbox: list[EmailMultiAlternatives],
 ) -> None:
     # Given
+    allowance = MAX_SEATS_IN_FREE_PLAN
+    usage = MAX_API_CALLS_IN_FREE_PLAN + 5_000
+
     now = timezone.now()
     assert organisation.is_paid is False
     assert organisation.subscription.is_free_plan is True
     assert organisation.subscription.max_api_calls == MAX_API_CALLS_IN_FREE_PLAN
 
+    OrganisationSubscriptionInformationCache.objects.create(
+        organisation=organisation,
+        allowed_30d_api_calls=allowance,
+        api_calls_30d=usage,
+    )
+
     mock_api_usage = mocker.patch(
         "organisations.task_helpers.get_current_api_usage",
     )
-    mock_api_usage.return_value = MAX_API_CALLS_IN_FREE_PLAN + 5_000
+    mock_api_usage.return_value = usage
 
     get_client_mock = mocker.patch("organisations.tasks.get_client")
     client_mock = MagicMock()
@@ -831,9 +848,6 @@ def test_handle_api_usage_notifications_missing_info_cache(
     organisation.subscription.plan = SCALE_UP
     organisation.subscription.save()
 
-    from organisations.task_helpers import logger
-
-    logger.addHandler(inspecting_handler)
     assert organisation.has_subscription_information_cache() is False
 
     mock_api_usage = mocker.patch(
@@ -860,10 +874,6 @@ def test_handle_api_usage_notifications_missing_info_cache(
         organisation=organisation,
     ).exists()
 
-    assert inspecting_handler.messages == [  # type: ignore[attr-defined]
-        f"Paid organisation {organisation.id} is missing subscription information cache"
-    ]
-
 
 @pytest.mark.parametrize("plan", ("scale-up", "scale-up-v2", "scale-up-v3"))
 @pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00")