From 351232fc82a88cd12483f701ba832d1ae3725dd5 Mon Sep 17 00:00:00 2001
From: Gagan Trivedi <gagandeeptrivedi47@gmail.com>
Date: Mon, 29 Jan 2024 16:23:29 +0530
Subject: [PATCH] feat(sse/tracking): Add project and org name to the influx
 event (#3337)

---
 api/sse/tasks.py                 | 40 ++++++++++++++++----------------
 api/tests/unit/sse/test_tasks.py | 30 ++++++++++++++++++++----
 2 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/api/sse/tasks.py b/api/sse/tasks.py
index 2bde04fed7e7..937d281b71e5 100644
--- a/api/sse/tasks.py
+++ b/api/sse/tasks.py
@@ -57,30 +57,30 @@ def update_sse_usage():
         ) as write_api:
             environments = Environment.objects.filter(
                 api_key__in=agg_request_count.keys()
-            ).values("api_key", "id", "project_id", "project__organisation_id")
+            ).values(
+                "api_key",
+                "id",
+                "project_id",
+                "project__name",
+                "project__organisation_id",
+                "project__organisation__name",
+            )
 
             for environment in environments:
-                record = _get_influx_point(
-                    environment["id"],
-                    environment["project_id"],
-                    environment["project__organisation_id"],
-                    agg_request_count[environment["api_key"]],
-                    agg_last_event_generated_at[environment["api_key"]],
+                time = agg_last_event_generated_at[environment["api_key"]]
+                count = agg_request_count[environment["api_key"]]
+                record = (
+                    Point("sse_call")
+                    .field("request_count", count)
+                    .tag("environment_id", environment["id"])
+                    .tag("project_id", environment["project_id"])
+                    .tag("project", environment["project__name"])
+                    .tag("organisation_id", environment["project__organisation_id"])
+                    .tag("organisation", environment["project__organisation__name"])
+                    .time(time)
                 )
-                write_api.write(bucket=settings.INFLUXDB_BUCKET, record=record)
-
 
-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", organisation_id)
-        .tag("project_id", project_id)
-        .tag("environment_id", environment_id)
-        .time(time)
-    )
+                write_api.write(bucket=settings.INFLUXDB_BUCKET, record=record)
 
 
 def get_auth_header():
diff --git a/api/tests/unit/sse/test_tasks.py b/api/tests/unit/sse/test_tasks.py
index aa55790c7966..6d2127c2803d 100644
--- a/api/tests/unit/sse/test_tasks.py
+++ b/api/tests/unit/sse/test_tasks.py
@@ -121,10 +121,30 @@ def test_track_sse_usage(
         [
             call("sse_call"),
             call().field("request_count", 2),
-            call().field().tag("organisation_id", environment.project.organisation_id),
-            call().field().tag().tag("project_id", environment.project_id),
-            call().field().tag().tag().tag("environment_id", environment.id),
-            call().field().tag().tag().tag().time(second_access_log.generated_at),
+            call().field().tag("environment_id", environment.id),
+            call().field().tag().tag("project_id", environment.project.id),
+            call().field().tag().tag().tag("project", environment.project.name),
+            call()
+            .field()
+            .tag()
+            .tag()
+            .tag()
+            .tag("organisation_id", environment.project.organisation.id),
+            call()
+            .field()
+            .tag()
+            .tag()
+            .tag()
+            .tag()
+            .tag("organisation", environment.project.organisation.name),
+            call()
+            .field()
+            .tag()
+            .tag()
+            .tag()
+            .tag()
+            .tag()
+            .time(second_access_log.generated_at),
         ]
     )
 
@@ -136,5 +156,5 @@ def test_track_sse_usage(
     assert write_method.call_count == 1
     write_method.assert_called_once_with(
         bucket=influxdb_bucket,
-        record=mocked_influx_point().field().tag().tag().tag().time(),
+        record=mocked_influx_point().field().tag().tag().tag().tag().tag().time(),
     )