refac: use one query to fetch environment data

api/sse/tasks.py

@@ -55,26 +55,30 @@ def update_sse_usage():
         with influxdb_client.write_api(
             write_options=WriteOptions(batch_size=100, flush_interval=1000)
         ) as write_api:
-            for api_key, request_count in agg_request_count.items():
-                environment = Environment.get_from_cache(api_key)
-
-                if not environment:
-                    logger.warning("Invalid api_key %s", api_key)
-                    continue
+            environments = Environment.objects.filter(
+                api_key__in=agg_request_count.keys()
+            ).values("api_key", "id", "project_id", "project__organisation_id")
 
+            for environment in environments:
                 record = _get_influx_point(
-                    environment, request_count, agg_last_event_generated_at[api_key]
+                    environment["id"],
+                    environment["project_id"],
+                    environment["project__organisation_id"],
+                    agg_request_count[environment["api_key"]],
+                    agg_last_event_generated_at[environment["api_key"]],
                 )
                 write_api.write(bucket=settings.SSE_INFLUXDB_BUCKET, record=record)
 
 
-def _get_influx_point(environment: Environment, count: int, time: str) -> Point:
+def _get_influx_point(
+    environment_id: int, project_id: int, organisation_id: int, count: int, time: str
+) -> Point:
     return (
         Point("sse_call")
         .field("request_count", count)
-        .tag("organisation_id", environment.project.organisation_id)
-        .tag("project_id", environment.project_id)
-        .tag("environment_id", environment.id)
+        .tag("organisation_id", organisation_id)
+        .tag("project_id", project_id)
+        .tag("environment_id", environment_id)
         .time(time)
     )
 

api/tests/unit/sse/test_tasks.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from typing import Callable
 from unittest.mock import call
 
 import pytest
@@ -88,7 +89,10 @@ def test_auth_header_raises_exception_if_token_not_set(settings):
 
 
 def test_track_sse_usage(
-    mocker: MockerFixture, environment: Environment, settings: SettingsWrapper
+    mocker: MockerFixture,
+    environment: Environment,
+    django_assert_num_queries: Callable,
+    settings: SettingsWrapper,
 ):
     # Given - two valid logs
     first_access_log = SSEAccessLogs(datetime.now().isoformat(), environment.api_key)
@@ -108,7 +112,8 @@ def test_track_sse_usage(
     mocked_influx_point = mocker.patch("sse.tasks.Point")
 
     # When
-    update_sse_usage()
+    with django_assert_num_queries(1):
+        update_sse_usage()
 
     # Then
     # Point was generated correctly