Add logic for BA endpoint to accept storage_path (#892)

Author: adrian-codecov

upload/tests/views/test_bundle_analysis.py

@@ -103,6 +103,97 @@ def test_upload_bundle_analysis_success(db, client, mocker, mock_redis):
     )
 
 
+@pytest.mark.django_db(databases={"default", "timeseries"})
+@override_settings(SHELTER_SHARED_SECRET="shelter-shared-secret")
+def test_upload_bundle_analysis_success_shelter(db, client, mocker, mock_redis):
+    upload = mocker.patch.object(TaskService, "upload")
+    mock_sentry_metrics = mocker.patch(
+        "upload.views.bundle_analysis.sentry_metrics.incr"
+    )
+    create_presigned_put = mocker.patch(
+        "services.archive.StorageService.create_presigned_put",
+        return_value="test-presigned-put",
+    )
+
+    repository = RepositoryFactory.create()
+    commit_sha = "6fd5b89357fc8cdf34d6197549ac7c6d7e5977ef"
+
+    client = APIClient()
+    client.credentials(HTTP_AUTHORIZATION=f"token {repository.upload_token}")
+
+    res = client.post(
+        reverse("upload-bundle-analysis"),
+        {
+            "commit": commit_sha,
+            "slug": f"{repository.author.username}::::{repository.name}",
+            "build": "test-build",
+            "buildURL": "test-build-url",
+            "job": "test-job",
+            "service": "test-service",
+            "compareSha": "6fd5b89357fc8cdf34d6197549ac7c6d7e5aaaaa",
+            "storage_path": "shelter/test/path.txt",
+            "upload_external_id": "test-47078f85-2cee-4511-b38d-183c334ef43b",
+        },
+        format="json",
+        headers={"User-Agent": "codecov-cli/0.4.7"},
+    )
+    assert res.status_code == 201
+
+    # returns presigned storage URL
+    assert res.json() == {"url": "test-presigned-put"}
+
+    create_presigned_put.assert_called_once_with("bundle-analysis", ANY, 30)
+    call = create_presigned_put.mock_calls[0]
+    _, storage_path, _ = call.args
+    match = re.match(r"v1/uploads/([\d\w\-]+)\.json", storage_path)
+    assert match
+    (reportid,) = match.groups()
+
+    # creates commit
+    commit = Commit.objects.get(commitid=commit_sha)
+    assert commit
+
+    # saves args in Redis
+    redis = get_redis_connection()
+    args = redis.rpop(f"uploads/{repository.repoid}/{commit_sha}/bundle_analysis")
+    assert json.loads(args) == {
+        "reportid": reportid,
+        "build": "test-build",
+        "build_url": "test-build-url",
+        "job": "test-job",
+        "service": "test-service",
+        "url": f"v1/uploads/{reportid}.json",
+        "commit": commit_sha,
+        "report_code": None,
+        "bundle_analysis_compare_sha": "6fd5b89357fc8cdf34d6197549ac7c6d7e5aaaaa",
+    }
+
+    # sets latest upload timestamp
+    ts = redis.get(f"latest_upload/{repository.repoid}/{commit_sha}/bundle_analysis")
+    assert ts
+
+    # triggers upload task
+    upload.assert_called_with(
+        commitid=commit_sha,
+        repoid=repository.repoid,
+        countdown=4,
+        report_code=None,
+        report_type="bundle_analysis",
+    )
+    mock_sentry_metrics.assert_called_with(
+        "upload",
+        tags={
+            "agent": "cli",
+            "version": "0.4.7",
+            "action": "bundle_analysis",
+            "endpoint": "bundle_analysis",
+            "repo_visibility": "private",
+            "is_using_shelter": "no",
+            "position": "end",
+        },
+    )
+
+
 @pytest.mark.django_db(databases={"default", "timeseries"})
 def test_upload_bundle_analysis_org_token(db, client, mocker, mock_redis):
     mocker.patch.object(TaskService, "upload")

upload/views/bundle_analysis.py

@@ -49,6 +49,8 @@ class UploadSerializer(serializers.Serializer):
     branch = serializers.CharField(required=False, allow_null=True)
     compareSha = serializers.CharField(required=False, allow_null=True)
     git_service = serializers.CharField(required=False, allow_null=True)
+    storage_path = serializers.CharField(required=False, allow_null=True)
+    upload_external_id = serializers.CharField(required=False, allow_null=True)
 
 
 class BundleAnalysisView(APIView, ShelterMixin):
@@ -116,12 +118,16 @@ def post(self, request: HttpRequest) -> Response:
             },
         )
 
-        upload_external_id = str(uuid.uuid4())
-        storage_path = StoragePaths.upload.path(upload_key=upload_external_id)
-        archive_service = ArchiveService(repo)
-        url = archive_service.storage.create_presigned_put(
-            get_bucket_name(), storage_path, 30
-        )
+        storage_path = data.get("storage_path", None)
+        upload_external_id = data.get("upload_external_id", None)
+        url = None
+        if not self.is_shelter_request():
+            upload_external_id = str(uuid.uuid4())
+            storage_path = StoragePaths.upload.path(upload_key=upload_external_id)
+            archive_service = ArchiveService(repo)
+            url = archive_service.storage.create_presigned_put(
+                get_bucket_name(), storage_path, 30
+            )
 
         task_arguments = {
             # these are used in the upload task when saving an upload record