fix(API Usage): Allow self-hosted API usage alerts (#6975)

Author: emyller

api/app_analytics/analytics_db_service.py

@@ -1,4 +1,4 @@
-from datetime import date, datetime, timedelta
+from datetime import datetime, timedelta
 
 import structlog
 from common.core.utils import using_database_replica
@@ -122,20 +122,31 @@ def get_usage_data_from_local_db(
     return map_annotated_api_usage_buckets_to_usage_data(qs)
 
 
-def get_total_events_count(organisation) -> int:  # type: ignore[no-untyped-def]
+def get_total_events_count(
+    organisation: Organisation,
+    date_start: datetime | None = None,
+    date_stop: datetime | None = None,
+) -> int:
     """
-    Return total number of events for an organisation in the last 30 days
+    Return total number of events for an organisation for a range, or last 30 days
     """
+    today = timezone.now().replace(hour=0, minute=0, second=0, microsecond=0)
+    date_start = date_start or (today - timedelta(days=30))
+    date_stop = date_stop or today
     if settings.USE_POSTGRES_FOR_ANALYTICS:
-        count = APIUsageBucket.objects.filter(
+        count: int = APIUsageBucket.objects.filter(
             environment_id__in=_get_environment_ids_for_org(organisation),
-            created_at__date__lte=date.today(),
-            created_at__date__gt=date.today() - timedelta(days=30),
+            created_at__date__lte=date_stop,
+            created_at__date__gt=date_start,
             bucket_size=constants.ANALYTICS_READ_BUCKET_SIZE,
         ).aggregate(total_count=Sum("total_count"))["total_count"]
     else:
-        count = get_events_for_organisation(organisation.id)
-    return count  # type: ignore[no-any-return]
+        count = get_events_for_organisation(
+            organisation.id,
+            date_start=date_start,
+            date_stop=date_stop,
+        )
+    return count
 
 
 def get_feature_evaluation_data(

api/integrations/flagsmith/data/environment.json

@@ -104,6 +104,19 @@
       "feature_state_value": null,
       "featurestate_uuid": "f976df2f-2341-4623-8425-d6eda23a2ebc",
       "multivariate_feature_state_values": []
+    },
+    {
+      "django_id": 1229327,
+      "enabled": false,
+      "feature": {
+        "id": 195393,
+        "name": "get_current_api_usage_deprecated",
+        "type": "STANDARD"
+      },
+      "feature_segment": null,
+      "feature_state_value": null,
+      "featurestate_uuid": "e7ed0d54-b17c-4df1-ab98-5f8dc9597127",
+      "multivariate_feature_state_values": []
     }
   ],
   "id": 0,

api/organisations/task_helpers.py

@@ -7,8 +7,10 @@
 from django.template.loader import render_to_string
 from django.utils import timezone
 
+from app_analytics.analytics_db_service import get_total_events_count
 from app_analytics.influxdb_wrapper import get_current_api_usage
 from core.helpers import get_current_site_url
+from integrations.flagsmith.client import get_client
 from organisations.models import (
     Organisation,
     OrganisationAPIUsageNotification,
@@ -126,7 +128,16 @@ def handle_api_usage_notification_for_organisation(organisation: Organisation) -
 
         allowed_api_calls = subscription_cache.allowed_30d_api_calls
 
-    api_usage = get_current_api_usage(organisation.id, period_starts_at)
+    flagsmith_client = get_client("local", local_eval=True)
+    flags = flagsmith_client.get_identity_flags(
+        organisation.flagsmith_identifier,
+        traits=organisation.flagsmith_on_flagsmith_api_traits,
+    )
+    # TODO: Default to get_total_events_count — https://github.com/Flagsmith/flagsmith/issues/6985
+    if flags.is_feature_enabled("get_current_api_usage_deprecated"):  # pragma: no cover
+        api_usage = get_total_events_count(organisation, period_starts_at)
+    else:
+        api_usage = get_current_api_usage(organisation.id, period_starts_at)
 
     # For some reason the allowed API calls is set to 0 so default to the max free plan.
     allowed_api_calls = allowed_api_calls or MAX_API_CALLS_IN_FREE_PLAN

api/organisations/tasks.py

@@ -12,6 +12,7 @@
     register_task_handler,
 )
 
+from app_analytics.analytics_db_service import get_total_events_count
 from app_analytics.influxdb_wrapper import get_current_api_usage
 from integrations.flagsmith.client import get_client
 from organisations import subscription_info_cache
@@ -214,10 +215,19 @@ def charge_for_api_call_count_overages():  # type: ignore[no-untyped-def]
             continue
 
         subscription_cache = organisation.subscription_information_cache
-        api_usage = get_current_api_usage(
-            organisation_id=organisation.id,
-            date_start=subscription_cache.current_billing_term_starts_at,
-        )
+        # TODO: Default to get_total_events_count — https://github.com/Flagsmith/flagsmith/issues/6985
+        if flags.is_feature_enabled(
+            "get_current_api_usage_deprecated"
+        ):  # pragma: no cover
+            api_usage = get_total_events_count(
+                organisation,
+                date_start=subscription_cache.current_billing_term_starts_at,
+            )
+        else:
+            api_usage = get_current_api_usage(
+                organisation_id=organisation.id,
+                date_start=subscription_cache.current_billing_term_starts_at,
+            )
 
         # Grace period for organisations < 200% of usage.
         if (
@@ -347,10 +357,19 @@ def restrict_use_due_to_api_limit_grace_period_over() -> None:
         OrganisationBreachedGracePeriod.objects.get_or_create(organisation=organisation)
 
         subscription_cache = organisation.subscription_information_cache
-        api_usage = get_current_api_usage(
-            organisation_id=organisation.id,
-            date_start=now - timedelta(days=30),
-        )
+        # TODO: Default to get_total_events_count — https://github.com/Flagsmith/flagsmith/issues/6985
+        if flags.is_feature_enabled(
+            "get_current_api_usage_deprecated"
+        ):  # pragma: no cover
+            api_usage = get_total_events_count(
+                organisation,
+                date_start=now - timedelta(days=30),
+            )
+        else:
+            api_usage = get_current_api_usage(
+                organisation_id=organisation.id,
+                date_start=now - timedelta(days=30),
+            )
         if api_usage / subscription_cache.allowed_30d_api_calls < 1.0:
             logger.info(
                 f"API use for organisation {organisation.id} has fallen to below limit, so not restricting use."

api/tests/unit/app_analytics/test_analytics_db_service.py

@@ -45,6 +45,7 @@ def cache(organisation: Organisation) -> OrganisationSubscriptionInformationCach
 
 
 @pytest.mark.use_analytics_db
+@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00")
 def test_get_usage_data_from_local_db__multiple_buckets__returns_aggregated_daily_data(  # type: ignore[no-untyped-def]
     organisation, environment, settings
 ):
@@ -197,6 +198,7 @@ def test_get_usage_data_from_local_db__environment_filter__returns_expected(
 
 
 @pytest.mark.use_analytics_db
+@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00")
 def test_get_usage_data_from_local_db__labels_filter__returns_expected(
     organisation: Organisation,
     environment: Environment,
@@ -269,6 +271,7 @@ def test_get_usage_data_from_local_db__labels_filter__returns_expected(
 
 
 @pytest.mark.use_analytics_db
+@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00")
 def test_get_total_events_count__multiple_buckets__returns_correct_total(  # type: ignore[no-untyped-def]
     organisation, environment, settings
 ):
@@ -322,6 +325,7 @@ def test_get_total_events_count__multiple_buckets__returns_correct_total(  # typ
 
 
 @pytest.mark.use_analytics_db
+@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00")
 def test_get_feature_evaluation_data_from_local_db__multiple_buckets__returns_aggregated_daily_data(
     feature: Feature,
     environment: Environment,
@@ -538,6 +542,7 @@ def test_get_usage_data__no_analytics_configured__no_calls_expected(
     mocked_get_usage_data_from_local_db.assert_not_called()
 
 
+@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00")
 def test_get_total_events_count__postgres_not_configured__calls_influx(  # type: ignore[no-untyped-def]
     mocker, settings, organisation
 ):
@@ -553,7 +558,9 @@ def test_get_total_events_count__postgres_not_configured__calls_influx(  # type:
     # Then
     assert total_events_count == mocked_get_events_for_organisation.return_value
     mocked_get_events_for_organisation.assert_called_once_with(
-        organisation_id=organisation.id
+        organisation.id,
+        date_start=datetime(2022, 12, 20, 0, 0, tzinfo=UTC),
+        date_stop=datetime(2023, 1, 19, 0, 0, tzinfo=UTC),
     )