fix task processor args

Author: khvn26

api/app_analytics/cache.py

@@ -1,9 +1,11 @@
-from collections import defaultdict
 from threading import Lock
 
 from django.conf import settings
 from django.utils import timezone
 
+from app_analytics.mappers import (
+    map_feature_evaluation_cache_to_track_feature_evaluations_by_environment_kwargs,
+)
 from app_analytics.models import Resource
 from app_analytics.tasks import (
     track_feature_evaluations_by_environment,
@@ -12,7 +14,6 @@
 from app_analytics.types import (
     APIUsageCacheKey,
     FeatureEvaluationCacheKey,
-    FeatureEvaluationKey,
     Labels,
 )
 
@@ -69,25 +70,10 @@ def __init__(self) -> None:
         self._lock = Lock()
 
     def _flush(self) -> None:
-        evaluation_data: dict[int, dict[FeatureEvaluationKey, int]] = defaultdict(dict)
-        for (
-            cache_key,
-            eval_count,
-        ) in self._cache.items():
-            key = FeatureEvaluationKey(
-                feature_name=cache_key.feature_name,
-                labels=cache_key.labels,
-            )
-            evaluation_data[cache_key.environment_id][key] = eval_count
-
-        # Schedule evaluation tracking by environment
-        for environment_id, feature_evaluations in evaluation_data.items():
-            track_feature_evaluations_by_environment.delay(
-                kwargs={
-                    "environment_id": environment_id,
-                    "feature_evaluations": list(feature_evaluations.items()),
-                }
-            )
+        for kwargs in map_feature_evaluation_cache_to_track_feature_evaluations_by_environment_kwargs(
+            self._cache
+        ):
+            track_feature_evaluations_by_environment.delay(kwargs=dict(kwargs))
 
         self._cache = {}
         self._last_flushed_at = timezone.now()

api/app_analytics/mappers.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from typing import Any, Iterable
 
 from django.http import HttpRequest
@@ -7,11 +8,14 @@
 
 from app_analytics.constants import TRACK_HEADERS
 from app_analytics.dataclasses import FeatureEvaluationData, UsageData
-from app_analytics.models import Resource
+from app_analytics.models import FeatureEvaluationRaw, Resource
 from app_analytics.types import (
     AnnotatedAPIUsageBucket,
     AnnotatedAPIUsageKey,
+    FeatureEvaluationCacheKey,
     Labels,
+    TrackFeatureEvaluationsByEnvironmentData,
+    TrackFeatureEvaluationsByEnvironmentKwargs,
 )
 
 _request_header_labels_model_fields: dict[str, Any] = {
@@ -90,3 +94,44 @@ def map_request_to_labels(request: HttpRequest) -> Labels:
         exclude_unset=True,
     )
     return result
+
+
+def map_feature_evaluation_cache_to_track_feature_evaluations_by_environment_kwargs(
+    cache: dict[FeatureEvaluationCacheKey, int],
+) -> list[TrackFeatureEvaluationsByEnvironmentKwargs]:
+    feature_evaluations_by_environment: dict[
+        int, list[TrackFeatureEvaluationsByEnvironmentData]
+    ] = defaultdict(list)
+
+    for cache_key, evaluation_count in cache.items():
+        environment_id = cache_key.environment_id
+        feature_evaluations_by_environment[environment_id].append(
+            {
+                "feature_name": cache_key.feature_name,
+                "labels": dict(cache_key.labels),
+                "evaluation_count": evaluation_count,
+            }
+        )
+
+    return [
+        {
+            "environment_id": environment_id,
+            "feature_evaluations": feature_evaluations,
+        }
+        for environment_id, feature_evaluations in feature_evaluations_by_environment.items()
+    ]
+
+
+def map_feature_evaluation_data_to_feature_evaluation_raw(
+    environment_id: int,
+    feature_evaluations: list[TrackFeatureEvaluationsByEnvironmentData],
+) -> list[FeatureEvaluationRaw]:
+    return [
+        FeatureEvaluationRaw(
+            environment_id=environment_id,
+            feature_name=feature_evaluation["feature_name"],
+            evaluation_count=feature_evaluation["evaluation_count"],
+            labels=feature_evaluation["labels"],
+        )
+        for feature_evaluation in feature_evaluations
+    ]

api/app_analytics/tasks.py

@@ -1,5 +1,5 @@
 from datetime import datetime, timedelta
-from typing import Any, List, Tuple
+from typing import Any, List, Tuple, Unpack
 
 from django.conf import settings
 from django.db.models import Q, Sum
@@ -11,6 +11,7 @@
 )
 
 from app_analytics.constants import ANALYTICS_READ_BUCKET_SIZE
+from app_analytics.mappers import map_feature_evaluation_data_to_feature_evaluation_raw
 from app_analytics.models import (
     APIUsageBucket,
     APIUsageRaw,
@@ -26,8 +27,8 @@
     track_feature_evaluation_influxdb_v2 as track_feature_evaluation_influxdb_v2_service,
 )
 from app_analytics.types import (
-    FeatureEvaluationKey,
     Labels,
+    TrackFeatureEvaluationsByEnvironmentKwargs,
 )
 from environments.models import Environment
 
@@ -95,21 +96,17 @@ def track_feature_evaluation_v2(
 
 @register_task_handler()
 def track_feature_evaluations_by_environment(
-    environment_id: int,
-    feature_evaluations: list[tuple[FeatureEvaluationKey, int]],
+    **kwargs: Unpack[TrackFeatureEvaluationsByEnvironmentKwargs],
 ) -> None:
+    environment_id = kwargs["environment_id"]
+    feature_evaluations = kwargs["feature_evaluations"]
     if settings.USE_POSTGRES_FOR_ANALYTICS:
-        feature_evaluation_objects = []
-        for (feature_name, labels), evaluation_count in feature_evaluations:
-            feature_evaluation_objects.append(
-                FeatureEvaluationRaw(
-                    feature_name=feature_name,
-                    environment_id=environment_id,
-                    evaluation_count=evaluation_count,
-                    labels=dict(labels),
-                )
+        FeatureEvaluationRaw.objects.bulk_create(
+            map_feature_evaluation_data_to_feature_evaluation_raw(
+                environment_id=environment_id,
+                feature_evaluations=feature_evaluations,
             )
-        FeatureEvaluationRaw.objects.bulk_create(feature_evaluation_objects)
+        )
     elif settings.INFLUXDB_TOKEN:
         track_feature_evaluation_influxdb(
             environment_id=environment_id,

api/app_analytics/track.py

@@ -8,7 +8,7 @@
 
 from app_analytics.influxdb_wrapper import InfluxDBWrapper
 from app_analytics.models import Resource
-from app_analytics.types import FeatureEvaluationKey, Labels
+from app_analytics.types import Label, Labels, TrackFeatureEvaluationsByEnvironmentData
 from environments.models import Environment
 from util.util import postpone
 
@@ -116,7 +116,7 @@ def track_request_influxdb(
 
 def track_feature_evaluation_influxdb(
     environment_id: int,
-    feature_evaluations: list[tuple[FeatureEvaluationKey, int]],
+    feature_evaluations: list[TrackFeatureEvaluationsByEnvironmentData],
 ) -> None:
     """
     Sends Feature analytics event data to InfluxDB
@@ -126,13 +126,20 @@ def track_feature_evaluation_influxdb(
     """
     influxdb = InfluxDBWrapper("feature_evaluation")  # type: ignore[no-untyped-call]
 
-    for (feature_name, labels), evaluation_count in feature_evaluations:
-        tags: dict[str, str | int] = {
-            "feature_id": feature_name,
+    for feature_evaluation in feature_evaluations:
+        tags: dict[str | Label, str | int] = {
+            "feature_id": feature_evaluation["feature_name"],
             "environment_id": environment_id,
-            **dict(labels),
+            **{
+                str(label): value
+                for label, value in feature_evaluation["labels"].items()
+            },
         }
-        influxdb.add_data_point("request_count", evaluation_count, tags=tags)
+        influxdb.add_data_point(
+            "request_count",
+            feature_evaluation["evaluation_count"],
+            tags=tags,
+        )
 
     influxdb.write()
 

api/app_analytics/types.py

@@ -24,9 +24,15 @@ class FeatureEvaluationCacheKey(NamedTuple):
     labels: tuple[tuple["Label", str], ...]
 
 
-class FeatureEvaluationKey(NamedTuple):
+class TrackFeatureEvaluationsByEnvironmentData(TypedDict):
     feature_name: str
-    labels: tuple[tuple["Label", str], ...]
+    labels: "Labels"
+    evaluation_count: int
+
+
+class TrackFeatureEvaluationsByEnvironmentKwargs(TypedDict):
+    environment_id: int
+    feature_evaluations: list[TrackFeatureEvaluationsByEnvironmentData]
 
 
 class AnnotatedAPIUsageBucket(TypedDict):

api/tests/unit/app_analytics/test_tasks.py

@@ -20,7 +20,7 @@
     track_feature_evaluations_by_environment,
     track_request,
 )
-from app_analytics.types import FeatureEvaluationKey
+from app_analytics.types import TrackFeatureEvaluationsByEnvironmentData
 from environments.models import Environment
 
 pytestmark = pytest.mark.skip_if_no_analytics_db
@@ -177,6 +177,7 @@ def test_track_request__postgres__inserts_expected(
 
 
 def test_track_request__influx__calls_expected(
+    db: None,
     settings: SettingsWrapper,
     mocker: MockerFixture,
     environment: Environment,
@@ -210,12 +211,23 @@ def test_track_feature_evaluation(settings: SettingsWrapper) -> None:
     settings.USE_POSTGRES_FOR_ANALYTICS = True
     environment_id = 1
     feature_evaluations = [
-        (FeatureEvaluationKey("feature1", ()), 10),
-        (FeatureEvaluationKey("feature2", ()), 20),
+        TrackFeatureEvaluationsByEnvironmentData(
+            feature_name="feature1",
+            labels={},
+            evaluation_count=10,
+        ),
+        TrackFeatureEvaluationsByEnvironmentData(
+            feature_name="feature2",
+            labels={},
+            evaluation_count=20,
+        ),
     ]
 
     # When
-    track_feature_evaluations_by_environment(environment_id, feature_evaluations)
+    track_feature_evaluations_by_environment(
+        environment_id=environment_id,
+        feature_evaluations=feature_evaluations,
+    )
 
     # Then
     assert (
@@ -245,12 +257,27 @@ def test_track_feature_evaluation__influx__calls_expected(
     )
     environment_id = 1
     feature_evaluations = [
-        (FeatureEvaluationKey("feature1", ()), 10),
-        (FeatureEvaluationKey("feature2", ()), 20),
+        (
+            TrackFeatureEvaluationsByEnvironmentData(
+                feature_name="feature1",
+                labels={},
+                evaluation_count=10,
+            )
+        ),
+        (
+            TrackFeatureEvaluationsByEnvironmentData(
+                feature_name="feature2",
+                labels={},
+                evaluation_count=20,
+            )
+        ),
     ]
 
     # When
-    track_feature_evaluations_by_environment(environment_id, feature_evaluations)
+    track_feature_evaluations_by_environment(
+        environment_id=environment_id,
+        feature_evaluations=feature_evaluations,
+    )
 
     # Then
     track_feature_evaluation_influxdb_mock.assert_called_once_with(

api/tests/unit/app_analytics/test_unit_app_analytics_cache.py

@@ -5,7 +5,7 @@
 
 from app_analytics.cache import APIUsageCache, FeatureEvaluationCache
 from app_analytics.models import Resource
-from app_analytics.types import FeatureEvaluationKey
+from app_analytics.types import TrackFeatureEvaluationsByEnvironmentData
 
 
 def test_api_usage_cache(
@@ -170,19 +170,15 @@ def test_feature_evaluation_cache(
                 kwargs={
                     "environment_id": 1,
                     "feature_evaluations": [
-                        (
-                            FeatureEvaluationKey(
-                                feature_name="feature_1_name",
-                                labels=(),
-                            ),
-                            11,
+                        TrackFeatureEvaluationsByEnvironmentData(
+                            feature_name="feature_1_name",
+                            labels={},
+                            evaluation_count=11,
                         ),
-                        (
-                            FeatureEvaluationKey(
-                                feature_name="feature_2_name",
-                                labels=(),
-                            ),
-                            10,
+                        TrackFeatureEvaluationsByEnvironmentData(
+                            feature_name="feature_2_name",
+                            labels={},
+                            evaluation_count=10,
                         ),
                     ],
                 }
@@ -191,12 +187,10 @@ def test_feature_evaluation_cache(
                 kwargs={
                     "environment_id": 2,
                     "feature_evaluations": [
-                        (
-                            FeatureEvaluationKey(
-                                feature_name="feature_2_name",
-                                labels=(),
-                            ),
-                            10,
+                        TrackFeatureEvaluationsByEnvironmentData(
+                            feature_name="feature_2_name",
+                            labels={},
+                            evaluation_count=10,
                         )
                     ],
                 }
@@ -205,19 +199,15 @@ def test_feature_evaluation_cache(
                 kwargs={
                     "environment_id": 1,
                     "feature_evaluations": [
-                        (
-                            FeatureEvaluationKey(
-                                feature_name="feature_1_name",
-                                labels=(("client_application_name", "test-app"),),
-                            ),
-                            1,
+                        TrackFeatureEvaluationsByEnvironmentData(
+                            feature_name="feature_1_name",
+                            labels={"client_application_name": "test-app"},
+                            evaluation_count=1,
                         ),
-                        (
-                            FeatureEvaluationKey(
-                                feature_name="feature_1_name",
-                                labels=(),
-                            ),
-                            1,
+                        TrackFeatureEvaluationsByEnvironmentData(
+                            feature_name="feature_1_name",
+                            labels={},
+                            evaluation_count=1,
                         ),
                     ],
                 }