fix(preprod): safely refresh size analysis processing state

src/sentry/preprod/api/endpoints/project_preprod_size.py

@@ -5,6 +5,7 @@
 
 import orjson
 import pydantic
+from django.utils import timezone
 from pydantic.tools import parse_obj_as
 from rest_framework import serializers
 from rest_framework.request import Request
@@ -98,12 +99,27 @@ def do_put(
             metrics.state = put.state
             metrics.error_code = put.error_code
             metrics.error_message = put.error_message
+            metrics.save()
         case PreprodArtifactSizeMetrics.SizeAnalysisState.PROCESSING:
-            metrics.state = put.state
+            # Reset the timeout and clear errors when an attempt starts.
+            # Check the state in the update so a concurrent completion is preserved.
+            PreprodArtifactSizeMetrics.objects.filter(
+                id=metrics.id,
+                state__in=[
+                    PreprodArtifactSizeMetrics.SizeAnalysisState.PENDING,
+                    PreprodArtifactSizeMetrics.SizeAnalysisState.PROCESSING,
+                    PreprodArtifactSizeMetrics.SizeAnalysisState.FAILED,
+                ],
+            ).update(
+                state=put.state,
+                error_code=None,
+                error_message=None,
+                date_updated=timezone.now(),
+            )
         case PreprodArtifactSizeMetrics.SizeAnalysisState.PENDING:
             metrics.state = put.state
+            metrics.save()
         case _:
             assert False, "unreachable"
-    metrics.save()
 
     return Response({"artifactId": str(head_artifact.id)})

tests/sentry/preprod/api/endpoints/test_project_preprod_size.py

@@ -1,8 +1,10 @@
+from datetime import timedelta
 from typing import cast
 
 import orjson
 from django.http import HttpResponse
 from django.test import override_settings
+from django.utils import timezone
 
 from sentry.preprod.models import PreprodArtifact, PreprodArtifactSizeMetrics
 from sentry.testutils.auth import generate_service_request_signature
@@ -160,12 +162,49 @@ def test_pending(self) -> None:
 
     @override_settings(LAUNCHPAD_RPC_SHARED_SECRET=[SHARED_SECRET_FOR_TESTS])
     def test_processing(self) -> None:
-        self.create_preprod_artifact_size_metrics(self.artifact)
-
-        self._put(orjson.dumps({"state": 1}))
-
-        metrics = PreprodArtifactSizeMetrics.objects.get(preprod_artifact=self.artifact)
-        assert metrics.state == PreprodArtifactSizeMetrics.SizeAnalysisState.PROCESSING
+        metrics = self.create_preprod_artifact_size_metrics(
+            self.artifact, state=PreprodArtifactSizeMetrics.SizeAnalysisState.PENDING
+        )
+        for state in [
+            PreprodArtifactSizeMetrics.SizeAnalysisState.PENDING,
+            PreprodArtifactSizeMetrics.SizeAnalysisState.PROCESSING,
+            PreprodArtifactSizeMetrics.SizeAnalysisState.FAILED,
+        ]:
+            with self.subTest(state=state):
+                old_time = timezone.now() - timedelta(minutes=35)
+                PreprodArtifactSizeMetrics.objects.filter(id=metrics.id).update(
+                    state=state,
+                    date_updated=old_time,
+                    error_code=PreprodArtifactSizeMetrics.ErrorCode.TIMEOUT,
+                    error_message="Previous attempt timed out",
+                )
+
+                response = self._put(orjson.dumps({"state": 1}))
+
+                assert response.status_code == 200
+                metrics.refresh_from_db()
+                assert metrics.state == PreprodArtifactSizeMetrics.SizeAnalysisState.PROCESSING
+                assert metrics.date_updated > old_time
+                assert metrics.error_code is None
+                assert metrics.error_message is None
+
+    @override_settings(LAUNCHPAD_RPC_SHARED_SECRET=[SHARED_SECRET_FOR_TESTS])
+    def test_processing_preserves_finished_results(self) -> None:
+        for state in [
+            PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED,
+            PreprodArtifactSizeMetrics.SizeAnalysisState.NOT_RAN,
+        ]:
+            with self.subTest(state=state):
+                self.artifact = self.create_preprod_artifact(project=self.project)
+                metrics = self.create_preprod_artifact_size_metrics(self.artifact, state=state)
+                original_updated = metrics.date_updated
+
+                response = self._put(orjson.dumps({"state": 1}))
+
+                assert response.status_code == 200
+                metrics.refresh_from_db()
+                assert metrics.state == state
+                assert metrics.date_updated == original_updated
 
     @override_settings(LAUNCHPAD_RPC_SHARED_SECRET=[SHARED_SECRET_FOR_TESTS])
     def test_requires_launchpad_rpc_authentication(self) -> None: