From 1cc91499417d7f92f50ee963931a75a7278e089a Mon Sep 17 00:00:00 2001
From: Gguidini <giovanni@codecov.io>
Date: Tue, 6 Feb 2024 14:39:18 -0300
Subject: [PATCH] fix: optimize n+1 query when creating upload

We noticed that we query the flags from an upload inside the loop
and decided to move that to the outside to reduce hits to the database
---
 upload/serializers.py            | 19 +++++++++++++------
 upload/tests/test_serializers.py | 19 +++++++++++++++++++
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/upload/serializers.py b/upload/serializers.py
index 3d8324542c..7a28406c73 100644
--- a/upload/serializers.py
+++ b/upload/serializers.py
@@ -56,6 +56,13 @@ def get_url(self, obj: ReportSession):
         commit = obj.report.commit
         return f"{settings.CODECOV_DASHBOARD_URL}/{repository.author.service}/{repository.author.username}/{repository.name}/commit/{commit.commitid}"
 
+    def _create_existing_flags_map(self, repoid: int) -> dict:
+        existing_flags = RepositoryFlag.objects.filter(repository_id=repoid).all()
+        existing_flags_map = {}
+        for flag_obj in existing_flags:
+            existing_flags_map[flag_obj.flag_name] = flag_obj
+        return existing_flags_map
+
     def create(self, validated_data):
         flag_names = (
             validated_data.pop("flags") if "flags" in validated_data.keys() else []
@@ -67,17 +74,17 @@ def create(self, validated_data):
         )
         upload = ReportSession.objects.create(**validated_data)
         flags = []
+
         if upload:
             repoid = upload.report.commit.repository.repoid
+            existing_flags_map = self._create_existing_flags_map(repoid)
             for individual_flag in flag_names:
-                existing_flag = RepositoryFlag.objects.filter(
-                    repository_id=repoid, flag_name=individual_flag
-                ).first()
-                if not existing_flag:
-                    existing_flag = RepositoryFlag.objects.create(
+                flag_obj = existing_flags_map.get(individual_flag, None)
+                if flag_obj is None:
+                    flag_obj = RepositoryFlag.objects.create(
                         repository_id=repoid, flag_name=individual_flag
                     )
-                flags.append(existing_flag)
+                flags.append(flag_obj)
             upload.flags.set(flags)
             return upload
 
diff --git a/upload/tests/test_serializers.py b/upload/tests/test_serializers.py
index 11aee60e2a..1f20f13edc 100644
--- a/upload/tests/test_serializers.py
+++ b/upload/tests/test_serializers.py
@@ -122,6 +122,25 @@ def test_upload_serializer_null_build_url_empty_flags(transactional_db, mocker):
     assert serializer.is_valid()
 
 
+def test__create_existing_flags_map(transactional_db, mocker):
+    mocker.patch(
+        "services.archive.StorageService.create_presigned_put",
+        return_value="presigned put",
+    )
+    upload = get_fake_upload_with_flags()
+    serializer = UploadSerializer(instance=upload)
+    flags_map = serializer._create_existing_flags_map(
+        upload.report.commit.repository.repoid
+    )
+    upload_flags = upload.flags.all()
+    flag1 = list(filter(lambda flag: flag.flag_name == "flag1", upload_flags))[0]
+    flag2 = list(filter(lambda flag: flag.flag_name == "flag2", upload_flags))[0]
+    assert flags_map == {
+        "flag1": flag1,
+        "flag2": flag2,
+    }
+
+
 def test_commit_serializer_contains_expected_fields(transactional_db, mocker):
     commit = CommitFactory.create()
     serializer = CommitSerializer(commit)