From 0ac7c40671f866b23151821ab9f0d8ab08666261 Mon Sep 17 00:00:00 2001
From: Matias Bordese <mbordese@gmail.com>
Date: Mon, 26 Aug 2024 12:03:26 -0300
Subject: [PATCH] Update org creation to use DB uniqueness constraint (#4926)

Fix issue related to
[logs](https://ops.grafana-ops.net/explore?schemaVersion=1&panes=%7B%2257p%22:%7B%22datasource%22:%22000000193%22,%22queries%22:%5B%7B%22refId%22:%22A%22,%22expr%22:%22%7Bnamespace%3D%5C%22amixr-prod%5C%22,%20cluster%3D%5C%22prod-us-central-0%5C%22,%20job%3D%5C%22amixr-prod%2Famixr-engine%5C%22%7D%20%7C%3D%20%5C%22user_management_organization.user_management_organization_stack_id_org_id_727b4929_uniq%5C%22%22,%22queryType%22:%22range%22,%22datasource%22:%7B%22type%22:%22loki%22,%22uid%22:%22000000193%22%7D,%22editorMode%22:%22code%22%7D%5D,%22range%22:%7B%22from%22:%22now-2d%22,%22to%22:%22now%22%7D%7D%7D&orgId=1)
(check for existing org using the unique DB index)
---
 engine/apps/grafana_plugin/helpers/gcom.py    | 22 +++++----
 engine/apps/grafana_plugin/tests/test_gcom.py | 47 +++++++++++++++++++
 2 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/engine/apps/grafana_plugin/helpers/gcom.py b/engine/apps/grafana_plugin/helpers/gcom.py
index e83ab9696a..d55c1b3652 100644
--- a/engine/apps/grafana_plugin/helpers/gcom.py
+++ b/engine/apps/grafana_plugin/helpers/gcom.py
@@ -60,18 +60,20 @@ def check_gcom_permission(token_string: str, context) -> GcomToken:
                 )
 
             # Get org from db or create a new one
-            organization, _ = Organization.objects.get_or_create(
+            organization, _ = Organization.objects.update_or_create(
                 stack_id=instance_info["id"],
-                stack_slug=instance_info["slug"],
-                grafana_url=instance_info["url"],
                 org_id=instance_info["orgId"],
-                org_slug=instance_info["orgSlug"],
-                org_title=instance_info["orgName"],
-                region_slug=instance_info["regionSlug"],
-                cluster_slug=instance_info["clusterSlug"],
-                gcom_token=token_string,
-                api_token=grafana_token,
-                defaults={"gcom_token_org_last_time_synced": timezone.now()},
+                defaults={
+                    "gcom_token_org_last_time_synced": timezone.now(),
+                    "stack_slug": instance_info["slug"],
+                    "grafana_url": instance_info["url"],
+                    "org_slug": instance_info["orgSlug"],
+                    "org_title": instance_info["orgName"],
+                    "region_slug": instance_info["regionSlug"],
+                    "cluster_slug": instance_info["clusterSlug"],
+                    "gcom_token": token_string,
+                    "api_token": grafana_token,
+                },
             )
     else:
         organization.stack_slug = instance_info["slug"]
diff --git a/engine/apps/grafana_plugin/tests/test_gcom.py b/engine/apps/grafana_plugin/tests/test_gcom.py
index c70f0719e4..b13ac0c130 100644
--- a/engine/apps/grafana_plugin/tests/test_gcom.py
+++ b/engine/apps/grafana_plugin/tests/test_gcom.py
@@ -3,6 +3,7 @@
 import pytest
 
 from apps.grafana_plugin.helpers.gcom import check_gcom_permission
+from apps.user_management.models import Organization
 
 
 @pytest.mark.parametrize(
@@ -59,3 +60,49 @@ def test_check_gcom_permission_updates_fields(make_organization, api_token, api_
     assert org.api_token == api_token if api_token_updated else broken_token
     assert org.gcom_token == gcom_token
     assert org.gcom_token_org_last_time_synced != last_time_gcom_synced
+
+
+@pytest.mark.django_db
+def test_check_gcom_permission_uniqueness_update_fields(make_organization):
+    gcom_token = "gcom:test_token"
+    fixed_token = "fixed_token"
+    instance_info = {
+        "id": 324534,
+        "slug": "testinstance",
+        "url": "http://example.com",
+        "orgId": 5671,
+        "orgSlug": "testorg",
+        "orgName": "Test Org",
+        "regionSlug": "us",
+        "clusterSlug": "us-test",
+    }
+    context = {
+        "stack_id": str(instance_info["id"]),
+        "org_id": str(instance_info["orgId"]),
+        "grafana_token": fixed_token,
+    }
+
+    org = make_organization(stack_id=instance_info["id"], org_id=instance_info["orgId"], api_token="broken_token")
+
+    # organization does not exist in the first check but it is created before the second check
+    with patch(
+        "apps.grafana_plugin.helpers.gcom.Organization.objects.filter", return_value=Organization.objects.none()
+    ):
+        with patch(
+            "apps.grafana_plugin.helpers.GcomAPIClient.get_instance_info",
+            return_value=instance_info,
+        ) as mock_instance_info:
+            check_gcom_permission(gcom_token, context)
+            mock_instance_info.assert_called()
+
+    org.refresh_from_db()
+    assert org.stack_id == instance_info["id"]
+    assert org.stack_slug == instance_info["slug"]
+    assert org.grafana_url == instance_info["url"]
+    assert org.org_id == instance_info["orgId"]
+    assert org.org_slug == instance_info["orgSlug"]
+    assert org.org_title == instance_info["orgName"]
+    assert org.region_slug == instance_info["regionSlug"]
+    assert org.cluster_slug == instance_info["clusterSlug"]
+    assert org.api_token == fixed_token
+    assert org.gcom_token == gcom_token