fix: API usage alerting in production

api/app_analytics/influxdb_wrapper.py

@@ -336,8 +336,7 @@ def get_current_api_usage(organisation_id: int, date_range: str) -> int:
         ),
         drop_columns=("_start", "_stop", "_time"),
         extra='|> sum() \
-              |> group() \
-              |> sort(columns: ["_value"], desc: true) ',
+               |> sort(columns: ["_value"], desc: true) ',
     )
 
     for result in results:
@@ -346,7 +345,7 @@ def get_current_api_usage(organisation_id: int, date_range: str) -> int:
             return 0
 
         # There should only be one matching result due to the
-        # group part of the query.
+        # sum part of the query.
         assert len(result.records) == 1
         return result.records[0].get_value()
 

api/integrations/flagsmith/client.py

@@ -8,11 +8,9 @@
 environment_flags = get_client().get_environment_flags()
 identity_flags = get_client().get_identity_flags()
 ```
-
-Possible extensions:
- - Allow for multiple clients?
 """
 import typing
+from time import sleep
 
 from django.conf import settings
 from flagsmith import Flagsmith
@@ -21,14 +19,24 @@
 from integrations.flagsmith.exceptions import FlagsmithIntegrationError
 from integrations.flagsmith.flagsmith_service import ENVIRONMENT_JSON_PATH
 
-_flagsmith_client: typing.Optional[Flagsmith] = None
+_flagsmith_clients: dict[str, Flagsmith] = {}
 
 
-def get_client() -> Flagsmith:
-    global _flagsmith_client
+def get_client(name: str = "default", local_eval: bool = False) -> Flagsmith:
+    global _flagsmith_clients
 
-    if not _flagsmith_client:
-        _flagsmith_client = Flagsmith(**_get_client_kwargs())
+    try:
+        _flagsmith_client = _flagsmith_clients[name]
+    except (KeyError, TypeError):
+        kwargs = _get_client_kwargs()
+        kwargs["enable_local_evaluation"] = local_eval
+        _flagsmith_client = Flagsmith(**kwargs)
+        # TODO: Remove this once client has been fixed.
+        # If the client is loaded then used too fast the flags
+        # returned from the method will be empty.
+        if local_eval:
+            sleep(0.5)
+        _flagsmith_clients[name] = _flagsmith_client
 
     return _flagsmith_client
 

api/organisations/tasks.py

@@ -8,6 +8,7 @@
 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 (
     OranisationAPIUsageNotification,
@@ -164,12 +165,20 @@ def _handle_api_usage_notifications(organisation: Organisation):
 
 
 def handle_api_usage_notifications():
+    flagsmith_client = get_client("local", local_eval=True)
+
     for organisation in Organisation.objects.filter(
         subscription_information_cache__current_billing_term_starts_at__isnull=False,
         subscription_information_cache__current_billing_term_ends_at__isnull=False,
     ).select_related(
         "subscription_information_cache",
     ):
+        if not flagsmith_client.get_identity_flags(
+            f"org.{organisation.id}.{organisation.name}",
+            traits={"organisation_id": organisation.id},
+        ).is_feature_enabled("api_usage_alerting"):
+            continue
+
         try:
             _handle_api_usage_notifications(organisation)
         except RuntimeError:

api/tests/unit/integrations/flagsmith/test_unit_flagsmith_client.py

@@ -12,7 +12,7 @@
 
 @pytest.fixture(autouse=True)
 def reset_globals(mocker: MockerFixture) -> None:
-    mocker.patch("integrations.flagsmith.client._flagsmith_client", None)
+    mocker.patch("integrations.flagsmith.client._flagsmith_clients", {})
     yield
 
 

api/tests/unit/organisations/test_unit_organisations_tasks.py

@@ -247,6 +247,8 @@ def test_handle_api_usage_notifications_below_100(
         "organisations.tasks.get_current_api_usage",
     )
     mock_api_usage.return_value = 91
+    mocker.patch("organisations.tasks.get_client")
+
     assert not OranisationAPIUsageNotification.objects.filter(
         organisation=organisation,
     ).exists()
@@ -334,6 +336,7 @@ def test_handle_api_usage_notifications_above_100(
     )
     mock_api_usage.return_value = 105
 
+    mocker.patch("organisations.tasks.get_client")
     assert not OranisationAPIUsageNotification.objects.filter(
         organisation=organisation,
     ).exists()