fix: API usage alerting in production (#3507)

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

api/app/pagination.py

@@ -4,7 +4,7 @@
 
 from drf_yasg import openapi
 from drf_yasg.inspectors import PaginatorInspector
-from flag_engine.identities.builders import build_identity_model
+from flag_engine.identities.models import IdentityModel
 from rest_framework.pagination import PageNumberPagination
 from rest_framework.response import Response
 
@@ -75,7 +75,7 @@ def paginate_queryset(self, dynamo_queryset, request, view=None):
             )
 
         return [
-            build_identity_model(identity_document)
+            IdentityModel.model_validate(identity_document)
             for identity_document in dynamo_queryset["Items"]
         ]
 

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/edge_api/identities/edge_identity_service.py

@@ -15,4 +15,9 @@ def get_edge_identity_overrides(
             environment_id=environment_id, feature_id=feature_id
         )
     )
-    return [IdentityOverrideV2.parse_obj(item) for item in override_items]
+    return [
+        IdentityOverrideV2.model_validate(
+            {**item, "environment_id": str(item["environment_id"])}
+        )
+        for item in override_items
+    ]

api/edge_api/identities/models.py

@@ -4,7 +4,6 @@
 
 from django.db.models import Prefetch, Q
 from flag_engine.features.models import FeatureStateModel
-from flag_engine.identities.builders import build_identity_model
 from flag_engine.identities.models import IdentityFeaturesList, IdentityModel
 
 from api_keys.models import MasterAPIKey
@@ -32,7 +31,7 @@ def __init__(self, engine_identity_model: IdentityModel):
 
     @classmethod
     def from_identity_document(cls, identity_document: dict) -> "EdgeIdentity":
-        return EdgeIdentity(build_identity_model(identity_document))
+        return EdgeIdentity(IdentityModel.model_validate(identity_document))
 
     @property
     def django_id(self) -> int:

api/edge_api/identities/views.py

@@ -7,7 +7,7 @@
 from django.shortcuts import get_object_or_404
 from django.utils.decorators import method_decorator
 from drf_yasg.utils import swagger_auto_schema
-from flag_engine.identities.builders import build_identity_model
+from flag_engine.identities.models import IdentityModel
 from flag_engine.identities.traits.models import TraitModel
 from pyngo import drf_error_details
 from rest_framework import status, viewsets
@@ -166,7 +166,7 @@ def perform_destroy(self, instance):
     )
     @action(detail=True, methods=["get"], url_path="list-traits")
     def get_traits(self, request, *args, **kwargs):
-        identity = build_identity_model(self.get_object())
+        identity = IdentityModel.model_validate(self.get_object())
         data = [trait.dict() for trait in identity.identity_traits]
         return Response(data=data, status=status.HTTP_200_OK)
 
@@ -180,7 +180,7 @@ def update_traits(self, request, *args, **kwargs):
         environment = self.get_environment_from_request()
         if not environment.project.organisation.persist_trait_data:
             raise TraitPersistenceError()
-        identity = build_identity_model(self.get_object())
+        identity = IdentityModel.model_validate(self.get_object())
         try:
             trait = TraitModel(**request.data)
         except pydantic.ValidationError as validation_error:

api/environments/dynamodb/wrappers/identity_wrapper.py

@@ -6,8 +6,7 @@
 from boto3.dynamodb.conditions import Key
 from django.conf import settings
 from django.core.exceptions import ObjectDoesNotExist
-from flag_engine.environments.builders import build_environment_model
-from flag_engine.identities.builders import build_identity_model
+from flag_engine.environments.models import EnvironmentModel
 from flag_engine.identities.models import IdentityModel
 from flag_engine.segments.evaluator import get_identity_segments
 from rest_framework.exceptions import NotFound
@@ -153,11 +152,11 @@ def get_segment_ids(
             raise ValueError("Must provide one of identity_pk or identity_model.")
 
         with suppress(ObjectDoesNotExist):
-            identity = identity_model or build_identity_model(
+            identity = identity_model or IdentityModel.model_validate(
                 self.get_item_from_uuid(identity_pk)
             )
             environment_wrapper = DynamoEnvironmentWrapper()
-            environment = build_environment_model(
+            environment = EnvironmentModel.model_validate(
                 environment_wrapper.get_item(identity.environment_api_key)
             )
             segments = get_identity_segments(environment, identity)

api/integrations/flagsmith/client.py

@@ -8,9 +8,6 @@
 environment_flags = get_client().get_environment_flags()
 identity_flags = get_client().get_identity_flags()
 ```
-
-Possible extensions:
- - Allow for multiple clients?
 """
 
 import typing
@@ -22,14 +19,19 @@
 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)
+        _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,21 @@ 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",
     ):
+        feature_enabled = flagsmith_client.get_identity_flags(
+            f"org.{organisation.id}.{organisation.name}",
+            traits={"organisation_id": organisation.id},
+        ).is_feature_enabled("api_usage_alerting")
+        if not feature_enabled:
+            continue
+
         try:
             _handle_api_usage_notifications(organisation)
         except RuntimeError: