diff --git a/src/sentry/objectstore/__init__.py b/src/sentry/objectstore/__init__.py index d677be1d29d26e..ab6be10d2e08d2 100644 --- a/src/sentry/objectstore/__init__.py +++ b/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 diff --git a/src/sentry/objectstore/endpoints/organization.py b/src/sentry/objectstore/endpoints/organization.py index 5bb22190d1a37f..28b66741b4beeb 100644 --- a/src/sentry/objectstore/endpoints/organization.py +++ b/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, diff --git a/src/sentry/preprod/snapshots/tasks.py b/src/sentry/preprod/snapshots/tasks.py index 5ce11f0f8b4ab6..0968de6cb46ecf 100644 --- a/src/sentry/preprod/snapshots/tasks.py +++ b/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")