fix(objectstore): Set auth token in proxy endpoint, preprod task

src/sentry/objectstore/__init__.py

@@ -1,5 +1,6 @@
 import subprocess
 from datetime import timedelta
+from typing import TypedDict, Unpack
 from urllib.parse import urlparse, urlunparse
 
 import urllib3
@@ -19,7 +20,7 @@
 from sentry.utils import metrics as sentry_metrics
 from sentry.utils.env import in_test_environment
 
-__all__ = ["get_attachments_session", "parse_accept_encoding"]
+__all__ = ["get_attachments_session", "get_preprod_session", "parse_accept_encoding"]
 
 
 def default_attachment_retention() -> int:
@@ -65,17 +66,21 @@ def distribution(
 _PREPROD_USECASE = Usecase("preprod", expiration_policy=TimeToLive(timedelta(days=30)))
 
 
-def create_client() -> Client:
+def create_client(token_override: str | None = None) -> Client:
     from sentry import options as options_store
 
     options = options_store.get("objectstore.config")
 
     # Initialize the `TokenGenerator` if key parameters are found.
-    token_generator = None
+    token = token_override
     if signing_key_options := options.get("token_generator"):
         # We require the `kid` and `secret_key` keys be set, other options are optional
-        if signing_key_options.get("kid") and signing_key_options.get("secret_key"):
-            token_generator = TokenGenerator(
+        if (
+            token is None
+            and signing_key_options.get("kid")
+            and signing_key_options.get("secret_key")
+        ):
+            token = TokenGenerator(
                 **signing_key_options,
             )
 
@@ -90,17 +95,31 @@ def create_client() -> Client:
             # Workaround for 0.0.14's default read timeout. Can be removed with 0.0.15
             {"timeout": urllib3.Timeout(connect=0.1)},
         ),
-        token=token_generator,
+        token=token,
     )
 
 
-def get_client() -> Client:
+def default_client() -> Client:
     global _OBJECTSTORE_CLIENT
     if not _OBJECTSTORE_CLIENT:
         _OBJECTSTORE_CLIENT = create_client()
     return _OBJECTSTORE_CLIENT
 
 
+def _get_session(
+    usecase: Usecase, org: int, project: int | None, client: Client | None = None
+) -> Session:
+    if client is None:
+        client = default_client()
+
+    if project is None:
+        session = client.session(usecase, org=org)
+    else:
+        session = client.session(usecase, org=org, project=project)
+
+    return session
+
+
 def get_attachments_usecase() -> Usecase:
     global _ATTACHMENTS_USECASE
     if not _ATTACHMENTS_USECASE:
@@ -112,11 +131,11 @@ def get_attachments_usecase() -> Usecase:
 
 
 def get_attachments_session(org: int, project: int) -> Session:
-    return get_client().session(get_attachments_usecase(), org=org, project=project)
+    return _get_session(get_attachments_usecase(), org, project)
 
 
-def get_preprod_session(org: int, project: int) -> Session:
-    return get_client().session(_PREPROD_USECASE, org=org, project=project)
+def get_preprod_session(org: int, project: int | None, client: Client | None = None) -> Session:
+    return _get_session(_PREPROD_USECASE, org, project, client=client)
 
 
 _IS_SYMBOLICATOR_CONTAINER: bool | None = None

src/sentry/objectstore/endpoints/organization.py

@@ -28,7 +28,7 @@
 from sentry.api.bases import OrganizationEndpoint
 from sentry.api.bases.organization import OrganizationReleasePermission
 from sentry.models.organization import Organization
-from sentry.objectstore import parse_accept_encoding
+from sentry.objectstore import get_preprod_session, parse_accept_encoding
 
 
 @cell_silo_endpoint
@@ -58,33 +58,34 @@ def get(
     ) -> Response | StreamingHttpResponse:
         if response := self._check_flag(request, organization):
             return response
-        return self._proxy(request, path)
+        return self._proxy(request, path, organization)
 
     def put(
         self, request: Request, organization: Organization, path: str
     ) -> Response | StreamingHttpResponse:
         if response := self._check_flag(request, organization):
             return response
-        return self._proxy(request, path)
+        return self._proxy(request, path, organization)
 
     def post(
         self, request: Request, organization: Organization, path: str
     ) -> Response | StreamingHttpResponse:
         if response := self._check_flag(request, organization):
             return response
-        return self._proxy(request, path)
+        return self._proxy(request, path, organization)
 
     def delete(
         self, request: Request, organization: Organization, path: str
     ) -> Response | StreamingHttpResponse:
         if response := self._check_flag(request, organization):
             return response
-        return self._proxy(request, path)
+        return self._proxy(request, path, organization)
 
     def _proxy(
         self,
         request: Request,
         path: str,
+        organization: Organization,
     ) -> Response | StreamingHttpResponse:
         assert request.method
         target_url = get_target_url(path)
@@ -94,6 +95,14 @@ def _proxy(
         headers.pop("Content-Length", None)
         headers.pop("Transfer-Encoding", None)
 
+        # Create an Objectstore session to mint a token
+        session = get_preprod_session(organization.id, None)
+        if token := session.mint_token():
+            headers["Authorization"] = f"Bearer {token}"
+        else:
+            # Objectstore can't parse Sentry's auth headers, must pop it
+            headers.pop("Authorization", None)
+
         response = requests.request(
             request.method,
             url=target_url,

src/sentry/preprod/snapshots/tasks.py

@@ -11,7 +11,7 @@
 from pydantic import ValidationError
 from taskbroker_client.retry import Retry
 
-from sentry.objectstore import get_preprod_session
+from sentry.objectstore import create_client, get_preprod_session
 from sentry.preprod.models import PreprodArtifact
 from sentry.preprod.snapshots.image_diff.compare import compare_images_batch
 from sentry.preprod.snapshots.image_diff.odiff import OdiffServer
@@ -125,6 +125,7 @@ def compare_snapshots(
     org_id: int,
     head_artifact_id: int,
     base_artifact_id: int,
+    storage_token: str | None = None,
 ) -> None:
     task_start_time = timezone.now()
     logger.info(
@@ -219,7 +220,8 @@ def compare_snapshots(
         )
 
     try:
-        session = get_preprod_session(org_id, project_id)
+        client_with_token = create_client(token_override=storage_token)
+        session = get_preprod_session(org_id, project_id, client=client_with_token)
 
         head_manifest_key = (head_metrics.extras or {}).get("manifest_key")
         base_manifest_key = (base_metrics.extras or {}).get("manifest_key")