fix(viewer-context): Populate org after request resolution

src/sentry/utils/sdk.py

@@ -30,6 +30,7 @@
 from sentry.utils import metrics
 from sentry.utils.db import DjangoAtomicIntegration
 from sentry.utils.rust import RustInfoIntegration
+from sentry.viewer_context import set_viewer_context_organization
 
 # Can't import models in utils because utils should be the bottom of the food chain
 if TYPE_CHECKING:
@@ -683,6 +684,7 @@ def bind_organization_context(organization: Organization | RpcOrganization) -> N
     helper = settings.SENTRY_ORGANIZATION_CONTEXT_HELPER
 
     scope = sentry_sdk.get_isolation_scope()
+    set_viewer_context_organization(organization.id)
 
     # XXX(dcramer): this is duplicated in organizationContext.jsx on the frontend
     with sentry_sdk.start_span(op="other", name="bind_organization_context"):

src/sentry/viewer_context.py

@@ -103,6 +103,15 @@ def get_viewer_context() -> ViewerContext | None:
     return _viewer_context_var.get()
 
 
+def set_viewer_context_organization(organization_id: int) -> None:
+    """Update the current ``ViewerContext`` with a resolved organization id."""
+    ctx = get_viewer_context()
+    if ctx is None or ctx.organization_id == organization_id:
+        return
+
+    _viewer_context_var.set(dataclasses.replace(ctx, organization_id=organization_id))
+
+
 # ---------------------------------------------------------------------------
 # JWT encoding / decoding for cross-service propagation
 # ---------------------------------------------------------------------------

tests/sentry/api/bases/test_organization.py

@@ -42,6 +42,7 @@
 from sentry.users.services.user.serial import serialize_rpc_user
 from sentry.users.services.user.service import user_service
 from sentry.utils.security.orgauthtoken_token import hash_token
+from sentry.viewer_context import ViewerContext, get_viewer_context, viewer_context_scope
 
 
 class MockSuperUser:
@@ -377,6 +378,20 @@ def build_request(self, user=None, active_superuser=False, **params):
         return request
 
 
+class OrganizationEndpointViewerContextTest(BaseOrganizationEndpointTest):
+    def test_convert_args_enriches_viewer_context_with_organization(self) -> None:
+        request = drf_request_from_request(self.build_request(user=self.owner))
+        request._request.organization = None
+
+        with viewer_context_scope(ViewerContext(user_id=self.owner.id)):
+            self.endpoint.convert_args(request, self.org.slug)
+            ctx = get_viewer_context()
+
+        assert ctx is not None
+        assert ctx.user_id == self.owner.id
+        assert ctx.organization_id == self.org.id
+
+
 class GetProjectIdsTest(BaseOrganizationEndpointTest):
     def setUp(self) -> None:
         self.team_1 = self.create_team(organization=self.org)

tests/sentry/api/bases/test_project.py

@@ -1,13 +1,17 @@
+from django.contrib.sessions.backends.base import SessionBase
+from django.test import RequestFactory
 from rest_framework.views import APIView
 
-from sentry.api.bases.project import ProjectAndStaffPermission, ProjectPermission
+from sentry.api.bases.project import ProjectAndStaffPermission, ProjectEndpoint, ProjectPermission
+from sentry.auth.access import from_request
 from sentry.models.apitoken import ApiToken
 from sentry.models.project import Project
 from sentry.testutils.cases import TestCase
 from sentry.testutils.helpers import with_feature
 from sentry.testutils.requests import drf_request_from_request
 from sentry.users.models.user import User
 from sentry.users.services.user.serial import serialize_rpc_user
+from sentry.viewer_context import ViewerContext, get_viewer_context, viewer_context_scope
 
 
 class ProjectPermissionBase(TestCase):
@@ -34,6 +38,26 @@ def has_object_perm(
         )
 
 
+class ProjectEndpointViewerContextTest(TestCase):
+    def test_convert_args_enriches_viewer_context_with_organization(self) -> None:
+        endpoint = ProjectEndpoint()
+        raw_request = RequestFactory().get("/")
+        raw_request.session = SessionBase()
+        raw_request.user = self.user
+        raw_request.auth = None
+        raw_request.access = from_request(drf_request_from_request(raw_request), self.organization)
+        request = drf_request_from_request(raw_request)
+        request._request.organization = None
+
+        with viewer_context_scope(ViewerContext(user_id=self.user.id)):
+            endpoint.convert_args(request, self.organization.slug, self.project.slug)
+            ctx = get_viewer_context()
+
+        assert ctx is not None
+        assert ctx.user_id == self.user.id
+        assert ctx.organization_id == self.organization.id
+
+
 class ProjectPermissionTest(ProjectPermissionBase):
     def test_regular_user(self) -> None:
         user = self.create_user(is_superuser=False)

tests/sentry/api/bases/test_team.py

@@ -1,12 +1,16 @@
+from django.contrib.sessions.backends.base import SessionBase
+from django.test import RequestFactory
 from rest_framework.views import APIView
 
-from sentry.api.bases.team import TeamPermission
+from sentry.api.bases.team import TeamEndpoint, TeamPermission
+from sentry.auth.access import from_request
 from sentry.models.apitoken import ApiToken
 from sentry.models.team import Team
 from sentry.testutils.cases import TestCase
 from sentry.testutils.helpers import with_feature
 from sentry.testutils.requests import drf_request_from_request
 from sentry.users.models.user import User
+from sentry.viewer_context import ViewerContext, get_viewer_context, viewer_context_scope
 
 
 class TeamPermissionBase(TestCase):
@@ -33,6 +37,27 @@ def has_object_perm(
         )
 
 
+class TeamEndpointViewerContextTest(TeamPermissionBase):
+    def test_convert_args_enriches_viewer_context_with_organization(self) -> None:
+        endpoint = TeamEndpoint()
+        self.create_member(user=self.user, organization=self.org, role="member", teams=[self.team])
+        raw_request = RequestFactory().get("/")
+        raw_request.session = SessionBase()
+        raw_request.user = self.user
+        raw_request.auth = None
+        raw_request.access = from_request(drf_request_from_request(raw_request), self.org)
+        request = drf_request_from_request(raw_request)
+        request._request.organization = None
+
+        with viewer_context_scope(ViewerContext(user_id=self.user.id)):
+            endpoint.convert_args(request, self.org.slug, self.team.slug)
+            ctx = get_viewer_context()
+
+        assert ctx is not None
+        assert ctx.user_id == self.user.id
+        assert ctx.organization_id == self.org.id
+
+
 class TeamPermissionTest(TeamPermissionBase):
     def test_get_regular_user(self) -> None:
         user = self.create_user()

tests/sentry/issues/endpoints/test_group.py

@@ -2,13 +2,15 @@
 
 from rest_framework.views import APIView
 
+from sentry.auth.access import from_request
 from sentry.issues.endpoints.bases.group import GroupAiEndpoint, GroupAiPermission
 from sentry.models.apitoken import ApiToken
 from sentry.models.group import Group
 from sentry.testutils.cases import TestCase
 from sentry.testutils.helpers.options import override_options
 from sentry.testutils.requests import drf_request_from_request
 from sentry.users.models.user import User
+from sentry.viewer_context import ViewerContext, get_viewer_context, viewer_context_scope
 
 
 class GroupAiPermissionTest(TestCase):
@@ -125,3 +127,24 @@ def setUp(self) -> None:
     def test_permission_classes(self) -> None:
         assert hasattr(self.endpoint, "permission_classes")
         assert self.endpoint.permission_classes == (GroupAiPermission,)
+
+    def test_convert_args_enriches_viewer_context_with_organization(self) -> None:
+        user = self.create_user()
+        self.create_member(
+            user=user,
+            organization=self.project.organization,
+            role="member",
+            teams=[self.project.teams.first()],
+        )
+        request = self.make_request(user=user, method="GET")
+        request.access = from_request(drf_request_from_request(request), self.project.organization)
+        request = drf_request_from_request(request)
+        request._request.organization = None
+
+        with viewer_context_scope(ViewerContext(user_id=user.id)):
+            self.endpoint.convert_args(request, str(self.group.id), self.project.organization.slug)
+            ctx = get_viewer_context()
+
+        assert ctx is not None
+        assert ctx.user_id == user.id
+        assert ctx.organization_id == self.project.organization.id

tests/sentry/test_viewer_context.py

@@ -10,6 +10,7 @@
     ActorType,
     ViewerContext,
     get_viewer_context,
+    set_viewer_context_organization,
     viewer_context_scope,
 )
 
@@ -125,3 +126,22 @@ def test_context_propagating_executor_does_not_leak_across_submissions(self):
 
         assert inside is ctx
         assert outside is None
+
+    def test_set_organization_updates_current_context(self):
+        ctx = ViewerContext(user_id=1, actor_type=ActorType.USER)
+
+        with viewer_context_scope(ctx):
+            set_viewer_context_organization(42)
+            updated = get_viewer_context()
+
+            assert updated is not None
+            assert updated.organization_id == 42
+            assert updated.user_id == ctx.user_id
+            assert updated.actor_type is ctx.actor_type
+
+        assert get_viewer_context() is None
+
+    def test_set_organization_is_noop_without_context(self):
+        set_viewer_context_organization(42)
+
+        assert get_viewer_context() is None