Add batch insert fn for measurements (#466)

Author: adrianviquez

shared/upload/utils.py

@@ -1,6 +1,7 @@
 from datetime import timedelta
 from enum import Enum
 
+from django.db import transaction
 from django.db.models import Q
 from django.utils import timezone
 
@@ -35,6 +36,18 @@ def query_monthly_coverage_measurements(plan_service: PlanService) -> int:
     return queryset[:monthly_limit].count()
 
 
+def bulk_insert_coverage_measurements(
+    measurements: list[UserMeasurement],
+) -> list[UserMeasurement]:
+    """
+    This function takes measurements as input and bulk_creates them into the DB.
+    The atomic transaction ensures either all transactions are inserted or none
+    if there's an error
+    """
+    with transaction.atomic():
+        return UserMeasurement.objects.bulk_create(measurements)
+
+
 def insert_coverage_measurement(
     owner_id: int,
     repo_id: int,

tests/unit/upload/test_utils.py

@@ -14,6 +14,7 @@
 from shared.django_apps.user_measurements.models import UserMeasurement
 from shared.plan.service import PlanService
 from shared.upload.utils import (
+    bulk_insert_coverage_measurements,
     insert_coverage_measurement,
     query_monthly_coverage_measurements,
 )
@@ -145,3 +146,28 @@ def test_query_monthly_coverage_measurements_beyond_monthly_limit(
         )
         # 10 uploads total, max 3 returned
         assert monthly_measurements == 3
+
+    def test_bulk_insert_user_measurements(self):
+        owner = OwnerFactory()
+        measurements = []
+        for _ in range(5):
+            repo = RepositoryFactory.create(author=owner)
+            commit = CommitFactory.create(repository=repo)
+            report = CommitReportFactory.create(commit=commit)
+            upload = UploadFactory.create(report=report)
+            measurements.append(
+                UserMeasurement(
+                    owner_id=owner.ownerid,
+                    repo_id=repo.repoid,
+                    commit_id=commit.id,
+                    upload_id=upload.id,
+                    uploader_used="CLI",
+                    private_repo=repo.private,
+                    report_type=report.report_type,
+                )
+            )
+
+        inserted_measurements = bulk_insert_coverage_measurements(
+            measurements=measurements
+        )
+        assert len(inserted_measurements) == 5