diff --git a/api/app_analytics/migrations/0009_apiusagebucket_host.py b/api/app_analytics/migrations/0009_apiusagebucket_host.py new file mode 100644 index 000000000000..117f62c11a33 --- /dev/null +++ b/api/app_analytics/migrations/0009_apiusagebucket_host.py @@ -0,0 +1,37 @@ +# Generated by Django 5.2.17 on 2026-09-05 02:32 + +from datetime import timedelta + +from django.db import migrations, models +from django.utils import timezone + + +def delete_buckets_the_rollup_may_recompute(apps, schema_editor): # type: ignore[no-untyped-def] + # Buckets created before this migration have host "" and hold the full + # count for their 15 minute window. The bucketing task rebuilds the last + # hour of buckets on every run, one bucket per host now, so an old bucket + # in that range would get new per-host buckets added next to it and the + # same requests would be counted twice. Delete the old buckets the task + # can still reach; its next run rebuilds them from the raw data. + APIUsageBucket = apps.get_model("app_analytics", "APIUsageBucket") + APIUsageBucket.objects.using(schema_editor.connection.alias).filter( + created_at__gte=timezone.now() - timedelta(hours=2) + ).delete() + + +class Migration(migrations.Migration): + + dependencies = [ + ("app_analytics", "0008_labels_jsonb"), + ] + + operations = [ + migrations.AddField( + model_name="apiusagebucket", + name="host", + field=models.CharField(default="", max_length=255), + ), + migrations.RunPython( + delete_buckets_the_rollup_may_recompute, migrations.RunPython.noop + ), + ] diff --git a/api/app_analytics/models.py b/api/app_analytics/models.py index 869762161692..1ceb829ad309 100644 --- a/api/app_analytics/models.py +++ b/api/app_analytics/models.py @@ -94,10 +94,11 @@ def check_overlapping_buckets(self, filters): # type: ignore[no-untyped-def] class APIUsageBucket(AbstractBucket): resource = models.IntegerField(choices=Resource.choices) + host = models.CharField(max_length=255, default="") @hook(BEFORE_CREATE) def check_overlapping_buckets(self): # type: ignore[no-untyped-def] - filter = models.Q(resource=self.resource) + filter = models.Q(resource=self.resource, host=self.host) super().check_overlapping_buckets(filter) # type: ignore[no-untyped-call] diff --git a/api/app_analytics/tasks.py b/api/app_analytics/tasks.py index b0a5b4a1b7d2..6db7c387d85c 100644 --- a/api/app_analytics/tasks.py +++ b/api/app_analytics/tasks.py @@ -191,6 +191,7 @@ def populate_api_usage_bucket( defaults={"total_count": row["count"]}, environment_id=row["environment_id"], resource=row["resource"], + host=row["host"], bucket_size=bucket_size, created_at=bucket_start_time, labels=row["labels"], @@ -229,12 +230,12 @@ def _get_api_usage_source_data( if source_bucket_size: return ( APIUsageBucket.objects.filter(filters, bucket_size=source_bucket_size) - .values("environment_id", "resource", "labels") + .values("environment_id", "resource", "host", "labels") .annotate(count=Sum("total_count")) ) return ( APIUsageRaw.objects.filter(filters) - .values("environment_id", "resource", "labels") + .values("environment_id", "resource", "host", "labels") .annotate( count=Sum("count"), ) diff --git a/api/tests/unit/app_analytics/test_migrations.py b/api/tests/unit/app_analytics/test_migrations.py index 636785ef8447..f84e2e886bb5 100644 --- a/api/tests/unit/app_analytics/test_migrations.py +++ b/api/tests/unit/app_analytics/test_migrations.py @@ -1,5 +1,8 @@ +from datetime import timedelta + import pytest from django.db import connections +from django.utils import timezone from django_test_migrations.migrator import Migrator pytestmark = pytest.mark.use_analytics_db @@ -202,3 +205,41 @@ def test_0008_labels_jsonb__hstore_columns__converts_to_jsonb( .labels == expected_labels ) + + +def test_0009_apiusagebucket_host__recent_bucket__deleted_and_old_kept( + analytics_migrator: Migrator, +) -> None: + # Given buckets from before the host column existed, one old and one + # recent enough for the rollup to recompute + old_state = analytics_migrator.apply_initial_migration( + ("app_analytics", "0008_labels_jsonb"), + ) + APIUsageBucket = old_state.apps.get_model("app_analytics", "APIUsageBucket") + now = timezone.now() + old_bucket = APIUsageBucket.objects.using("analytics").create( + environment_id=1, + bucket_size=15, + created_at=now - timedelta(hours=3), + total_count=10, + resource=1, + ) + APIUsageBucket.objects.using("analytics").create( + environment_id=1, + bucket_size=15, + created_at=now - timedelta(minutes=30), + total_count=10, + resource=1, + ) + + # When + new_state = analytics_migrator.apply_tested_migration( + ("app_analytics", "0009_apiusagebucket_host"), + ) + + # Then only the old bucket remains, with an empty host + NewAPIUsageBucket = new_state.apps.get_model("app_analytics", "APIUsageBucket") + remaining = list( + NewAPIUsageBucket.objects.using("analytics").values_list("id", "host") + ) + assert remaining == [(old_bucket.id, "")] diff --git a/api/tests/unit/app_analytics/test_tasks.py b/api/tests/unit/app_analytics/test_tasks.py index be3db9212e59..24856c39464d 100644 --- a/api/tests/unit/app_analytics/test_tasks.py +++ b/api/tests/unit/app_analytics/test_tasks.py @@ -26,10 +26,12 @@ pytestmark = pytest.mark.use_analytics_db -def _create_api_usage_event(environment_id: int, when: datetime) -> APIUsageRaw: +def _create_api_usage_event( + environment_id: int, when: datetime, host: str = "host1" +) -> APIUsageRaw: event = APIUsageRaw.objects.create( environment_id=environment_id, - host="host1", + host=host, resource=Resource.FLAGS, ) # update created_at @@ -534,6 +536,58 @@ def test_populate_api_usage_bucket__source_bucket_size__aggregates_correctly( assert APIUsageBucket.objects.filter(bucket_size=15, total_count=300).count() == 1 +def test_populate_api_usage_bucket__multiple_hosts__preserves_host( + freezer: FrozenDateTimeFactory, +) -> None: + # Given events from two hosts in the same bucket window + environment_id = 1 + ninety_minutes_ago = timezone.now() - timedelta(minutes=90) + for _ in range(3): + _create_api_usage_event(environment_id, ninety_minutes_ago, host="edge-proxy") + _create_api_usage_event(environment_id, ninety_minutes_ago) + + # When + freezer.move_to(timezone.now() - timedelta(hours=1)) + populate_api_usage_bucket(bucket_size=15, run_every=60) + + # Then the buckets are split by host, each keeping its host + buckets = APIUsageBucket.objects.filter(environment_id=environment_id) + assert {(bucket.host, bucket.total_count) for bucket in buckets} == { + ("edge-proxy", 3), + ("host1", 1), + } + + +@pytest.mark.freeze_time("2023-01-19T09:00:00+00:00") +@pytest.mark.use_analytics_db +def test_populate_api_usage_bucket__source_buckets_with_hosts__preserves_host( + freezer: FrozenDateTimeFactory, +) -> None: + # Given two source buckets in the same window, from different hosts + environment_id = 1 + now = timezone.now() + for host, total_count in (("edge-proxy", 100), ("host1", 50)): + APIUsageBucket.objects.create( + environment_id=environment_id, + resource=Resource.FLAGS, + host=host, + total_count=total_count, + created_at=now, + bucket_size=5, + ) + freezer.move_to(timezone.now().replace(minute=47)) + + # When + populate_api_usage_bucket(bucket_size=15, run_every=60, source_bucket_size=5) + + # Then each host keeps its own bucket and total + buckets = APIUsageBucket.objects.filter(bucket_size=15) + assert {(bucket.host, bucket.total_count) for bucket in buckets} == { + ("edge-proxy", 100), + ("host1", 50), + } + + def _create_feature_evaluation_event( environment_id: int, feature_name: str,