feat: Backend for SDK metrics (#5623)

Co-authored-by: Matthew Elwell <matthew.elwell@flagsmith.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

api/app/settings/common.py

@@ -79,6 +79,7 @@
     "django.contrib.contenttypes",
     "django.contrib.sessions",
     "django.contrib.messages",
+    "django.contrib.postgres",
     "django.contrib.staticfiles",
     "django.contrib.humanize",
     "rest_framework",

api/app_analytics/analytics_db_service.py

@@ -1,12 +1,13 @@
 from datetime import date, datetime, timedelta
-from typing import TYPE_CHECKING, List
 
+import structlog
 from dateutil.relativedelta import relativedelta
 from django.conf import settings
-from django.db.models import Sum
+from django.db.models import Q, Sum
 from django.utils import timezone
 from rest_framework.exceptions import NotFound
 
+from app_analytics import constants
 from app_analytics.dataclasses import FeatureEvaluationData, UsageData
 from app_analytics.influxdb_wrapper import (
     get_events_for_organisation,
@@ -17,86 +18,60 @@
 from app_analytics.influxdb_wrapper import (
     get_usage_data as get_usage_data_from_influxdb,
 )
+from app_analytics.mappers import map_annotated_api_usage_buckets_to_usage_data
 from app_analytics.models import (
     APIUsageBucket,
     FeatureEvaluationBucket,
-    Resource,
 )
+from app_analytics.types import Labels, PeriodType
 from environments.models import Environment
 from features.models import Feature
-from organisations.models import Organisation
+from organisations.models import Organisation, OrganisationSubscriptionInformationCache
 
-from . import constants
-from .types import PERIOD_TYPE
+logger = structlog.get_logger("app_analytics")
 
 
 def get_usage_data(
     organisation: Organisation,
     environment_id: int | None = None,
     project_id: int | None = None,
-    period: PERIOD_TYPE | None = None,
+    period: PeriodType | None = None,
+    labels_filter: Labels | None = None,
 ) -> list[UsageData]:
-    now = timezone.now()
-    date_start = date_stop = None
-    sub_cache = getattr(organisation, "subscription_information_cache", None)
+    sub_cache = OrganisationSubscriptionInformationCache.objects.filter(
+        organisation=organisation
+    ).first()
 
-    is_subscription_valid = (
-        sub_cache is not None and sub_cache.is_billing_terms_dates_set()
+    date_start, date_stop = _get_start_date_and_stop_date_for_subscribed_organisation(
+        sub_cache=sub_cache,
+        period=period,
     )
 
-    if period in (constants.CURRENT_BILLING_PERIOD, constants.PREVIOUS_BILLING_PERIOD):
-        if not is_subscription_valid:
-            raise NotFound("No billing periods found for this organisation.")
-
-    if TYPE_CHECKING:
-        assert sub_cache is not None
-
-    match period:
-        case constants.CURRENT_BILLING_PERIOD:
-            starts_at = sub_cache.current_billing_term_starts_at or now - timedelta(
-                days=30
-            )
-            month_delta = relativedelta(now, starts_at).months
-            date_start = relativedelta(months=month_delta) + starts_at
-            date_stop = now
-
-        case constants.PREVIOUS_BILLING_PERIOD:
-            starts_at = sub_cache.current_billing_term_starts_at or now - timedelta(
-                days=30
-            )
-            month_delta = relativedelta(now, starts_at).months - 1
-            month_delta += relativedelta(now, starts_at).years * 12
-            date_start = relativedelta(months=month_delta) + starts_at
-            date_stop = relativedelta(months=month_delta + 1) + starts_at
-        case constants.NINETY_DAY_PERIOD:
-            date_start = now - relativedelta(days=90)
-            date_stop = now
-
     if settings.USE_POSTGRES_FOR_ANALYTICS:
-        kwargs = {
-            "organisation": organisation,
-            "environment_id": environment_id,
-            "project_id": project_id,
-        }
-        if date_start:
-            assert date_stop
-            kwargs["date_start"] = date_start  # type: ignore[assignment]
-            kwargs["date_stop"] = date_stop  # type: ignore[assignment]
-
-        return get_usage_data_from_local_db(**kwargs)  # type: ignore[arg-type]
-
-    kwargs = {
-        "organisation_id": organisation.id,
-        "environment_id": environment_id,
-        "project_id": project_id,
-    }
+        return get_usage_data_from_local_db(
+            organisation=organisation,
+            environment_id=environment_id,
+            project_id=project_id,
+            date_start=date_start,
+            date_stop=date_stop,
+            labels_filter=labels_filter,
+        )
 
-    if date_start:
-        assert date_stop
-        kwargs["date_start"] = date_start  # type: ignore[assignment]
-        kwargs["date_stop"] = date_stop  # type: ignore[assignment]
+    if settings.INFLUXDB_TOKEN:
+        return get_usage_data_from_influxdb(
+            organisation_id=organisation.id,
+            environment_id=environment_id,
+            project_id=project_id,
+            date_start=date_start,
+            date_stop=date_stop,
+            labels_filter=labels_filter,
+        )
 
-    return get_usage_data_from_influxdb(**kwargs)  # type: ignore[arg-type]
+    logger.warning(
+        "no-analytics-database-configured",
+        details=constants.NO_ANALYTICS_DATABASE_CONFIGURED_WARNING,
+    )
+    return []
 
 
 def get_usage_data_from_local_db(
@@ -105,7 +80,8 @@ def get_usage_data_from_local_db(
     project_id: int | None = None,
     date_start: datetime | None = None,
     date_stop: datetime | None = None,
-) -> List[UsageData]:
+    labels_filter: Labels | None = None,
+) -> list[UsageData]:
     if date_start is None:
         date_start = timezone.now() - timedelta(days=30)
     if date_stop is None:
@@ -127,28 +103,20 @@ def get_usage_data_from_local_db(
     if environment_id:
         qs = qs.filter(environment_id=environment_id)
 
+    if labels_filter:
+        qs = qs.filter(labels__contains=labels_filter)
+
     qs = (
         qs.filter(  # type: ignore[assignment]
             created_at__date__lte=date_stop,
             created_at__date__gt=date_start,
         )
         .order_by("created_at__date")
-        .values("created_at__date", "resource")
+        .values("created_at__date", "resource", "labels")
         .annotate(count=Sum("total_count"))
     )
-    data_by_day = {}
-    for row in qs:  # TODO Write proper mappers for this?
-        day = row["created_at__date"]
-        if day not in data_by_day:
-            data_by_day[day] = UsageData(day=day)
-        if column_name := Resource(row["resource"]).column_name:
-            setattr(
-                data_by_day[day],
-                column_name,
-                row["count"],
-            )
 
-    return data_by_day.values()  # type: ignore[return-value]
+    return map_annotated_api_usage_buckets_to_usage_data(qs)
 
 
 def get_total_events_count(organisation) -> int:  # type: ignore[no-untyped-def]
@@ -168,30 +136,53 @@ def get_total_events_count(organisation) -> int:  # type: ignore[no-untyped-def]
 
 
 def get_feature_evaluation_data(
-    feature: Feature, environment_id: int, period: int = 30
-) -> List[FeatureEvaluationData]:
+    feature: Feature,
+    environment_id: int,
+    period_days: int = 30,
+    labels_filter: Labels | None = None,
+) -> list[FeatureEvaluationData]:
     if settings.USE_POSTGRES_FOR_ANALYTICS:
         return get_feature_evaluation_data_from_local_db(
-            feature, environment_id, period
+            feature=feature,
+            environment_id=environment_id,
+            period_days=period_days,
+            labels_filter=labels_filter,
         )
-    return get_feature_evaluation_data_from_influxdb(
-        feature_name=feature.name, environment_id=environment_id, period=f"{period}d"
+
+    if settings.INFLUXDB_TOKEN:
+        return get_feature_evaluation_data_from_influxdb(
+            feature_name=feature.name,
+            environment_id=environment_id,
+            period_days=period_days,
+            labels_filter=labels_filter,
+        )
+
+    logger.warning(
+        "no-analytics-database-configured",
+        details=constants.NO_ANALYTICS_DATABASE_CONFIGURED_WARNING,
     )
+    return []
 
 
 def get_feature_evaluation_data_from_local_db(
-    feature: Feature, environment_id: int, period: int = 30
-) -> List[FeatureEvaluationData]:
+    feature: Feature,
+    environment_id: int,
+    period_days: int = 30,
+    labels_filter: Labels | None = None,
+) -> list[FeatureEvaluationData]:
+    filter = Q(
+        environment_id=environment_id,
+        bucket_size=constants.ANALYTICS_READ_BUCKET_SIZE,
+        feature_name=feature.name,
+        created_at__date__lte=timezone.now(),
+        created_at__date__gt=timezone.now() - timedelta(days=period_days),
+    )
+    if labels_filter:
+        filter &= Q(labels__contains=labels_filter)
     feature_evaluation_data = (
-        FeatureEvaluationBucket.objects.filter(
-            environment_id=environment_id,
-            bucket_size=constants.ANALYTICS_READ_BUCKET_SIZE,
-            feature_name=feature.name,
-            created_at__date__lte=timezone.now(),
-            created_at__date__gt=timezone.now() - timedelta(days=period),
-        )
+        FeatureEvaluationBucket.objects.filter(filter)
         .order_by("created_at__date")
-        .values("created_at__date", "feature_name", "environment_id")
+        .values("created_at__date", "feature_name", "environment_id", "labels")
         .annotate(count=Sum("total_count"))
     )
     usage_list = []
@@ -200,15 +191,61 @@ def get_feature_evaluation_data_from_local_db(
             FeatureEvaluationData(
                 day=data["created_at__date"],
                 count=data["count"],
+                labels=data["labels"],
             )
         )
     return usage_list
 
 
-def _get_environment_ids_for_org(organisation) -> List[int]:  # type: ignore[no-untyped-def]
+def _get_environment_ids_for_org(organisation: Organisation) -> list[int]:
     # We need to do this to prevent Django from generating a query that
     # references the environments and projects tables,
     # as they do not exist in the analytics database.
     return [
-        e.id for e in Environment.objects.filter(project__organisation=organisation)
+        *Environment.objects.filter(
+            project__organisation=organisation,
+        ).values_list(
+            "id",
+            flat=True,
+        )
     ]
+
+
+def _get_start_date_and_stop_date_for_subscribed_organisation(
+    sub_cache: OrganisationSubscriptionInformationCache | None,
+    period: PeriodType | None = None,
+) -> tuple[datetime | None, datetime | None]:
+    """
+    Populate start and stop date for the given period
+    from the organisation's subscription information.
+    """
+    now = timezone.now()
+
+    match period:
+        case constants.CURRENT_BILLING_PERIOD:
+            if sub_cache and sub_cache.current_billing_term_starts_at:
+                starts_at = sub_cache.current_billing_term_starts_at
+            else:
+                raise NotFound("No billing periods found for this organisation.")
+
+            month_delta = relativedelta(now, starts_at).months
+            date_start = relativedelta(months=month_delta) + starts_at
+            return date_start, now
+
+        case constants.PREVIOUS_BILLING_PERIOD:
+            if sub_cache and sub_cache.current_billing_term_starts_at:
+                starts_at = sub_cache.current_billing_term_starts_at
+            else:
+                raise NotFound("No billing periods found for this organisation.")
+
+            month_delta = relativedelta(now, starts_at).months - 1
+            month_delta += relativedelta(now, starts_at).years * 12
+            date_start = relativedelta(months=month_delta) + starts_at
+            date_stop = relativedelta(months=month_delta + 1) + starts_at
+            return date_start, date_stop
+
+        case constants.NINETY_DAY_PERIOD:
+            date_start = now - relativedelta(days=90)
+            return date_start, now
+
+    return None, None