From 5dafecf0d5ca612c2cc0caabfdd2c69e1bb71e5b Mon Sep 17 00:00:00 2001
From: Gagan <gagandeeptrivedi47@gmail.com>
Date: Mon, 17 Jun 2024 16:35:45 +0530
Subject: [PATCH] fix(postgres-analytics/usage): fix project_id filter (#4171)

---
 api/app_analytics/analytics_db_service.py     |  9 +++-
 .../test_analytics_db_service.py              | 52 +++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/api/app_analytics/analytics_db_service.py b/api/app_analytics/analytics_db_service.py
index 8c2d191eadc7..5ccd612247db 100644
--- a/api/app_analytics/analytics_db_service.py
+++ b/api/app_analytics/analytics_db_service.py
@@ -117,7 +117,14 @@ def get_usage_data_from_local_db(
         bucket_size=constants.ANALYTICS_READ_BUCKET_SIZE,
     )
     if project_id:
-        qs = qs.filter(project_id=project_id)
+        environment_ids = Environment.objects.filter(project_id=project_id).values_list(
+            "id", flat=True
+        )
+        # evaluate the queryset because analytics db does not have
+        # access to environment/project table
+        environment_ids = list(environment_ids)
+        qs = qs.filter(environment_id__in=environment_ids)
+
     if environment_id:
         qs = qs.filter(environment_id=environment_id)
 
diff --git a/api/tests/unit/app_analytics/test_analytics_db_service.py b/api/tests/unit/app_analytics/test_analytics_db_service.py
index 7941d3d1de4a..88e6fb634fd3 100644
--- a/api/tests/unit/app_analytics/test_analytics_db_service.py
+++ b/api/tests/unit/app_analytics/test_analytics_db_service.py
@@ -28,6 +28,7 @@
     Organisation,
     OrganisationSubscriptionInformationCache,
 )
+from projects.models import Project
 
 
 @pytest.fixture
@@ -105,6 +106,57 @@ def test_get_usage_data_from_local_db(organisation, environment, settings):
         assert data.day == today - timedelta(days=29 - count)
 
 
+@pytest.mark.skipif(
+    "analytics" not in settings.DATABASES,
+    reason="Skip test if analytics database is configured",
+)
+@pytest.mark.django_db(databases=["analytics", "default"])
+def test_get_usage_data_from_local_db_project_id_filter(
+    organisation: Organisation,
+    project: Project,
+    project_two: Project,
+    environment: Environment,
+    environment_two: Environment,
+    project_two_environment: Environment,
+    settings: SettingsWrapper,
+):
+    # Given
+    environment_id = environment.id
+    now = timezone.now()
+    read_bucket_size = 15
+    settings.ANALYTICS_BUCKET_SIZE = read_bucket_size
+    total_count = 10
+
+    # crate one bucket for every environment
+    for environment_id in [
+        environment.id,
+        environment_two.id,
+        project_two_environment.id,
+    ]:
+        APIUsageBucket.objects.create(
+            environment_id=environment_id,
+            resource=Resource.FLAGS,
+            total_count=total_count,
+            bucket_size=read_bucket_size,
+            created_at=now,
+        )
+    # When
+    usage_data_for_project_one = get_usage_data_from_local_db(
+        organisation, project_id=project.id
+    )
+    usage_data_for_project_two = get_usage_data_from_local_db(
+        organisation, project_id=project_two.id
+    )
+
+    # Then
+    assert len(usage_data_for_project_one) == 1
+    assert len(usage_data_for_project_two) == 1
+    assert (
+        list(usage_data_for_project_one)[0].flags == total_count * 2
+    )  # 2 environments
+    assert list(usage_data_for_project_two)[0].flags == total_count  # 1 environment
+
+
 @pytest.mark.skipif(
     "analytics" not in settings.DATABASES,
     reason="Skip test if analytics database is configured",