diff --git a/src/sentry/utils/sdk.py b/src/sentry/utils/sdk.py index e825d54e31dc31..e549cfd8b62dc2 100644 --- a/src/sentry/utils/sdk.py +++ b/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"): diff --git a/src/sentry/viewer_context.py b/src/sentry/viewer_context.py index 02029dd2c3700b..80bed8c7a21fc5 100644 --- a/src/sentry/viewer_context.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/tests/sentry/api/bases/test_organization.py b/tests/sentry/api/bases/test_organization.py index 5fbaed64c9002e..2cffd5a8631f87 100644 --- a/tests/sentry/api/bases/test_organization.py +++ b/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) diff --git a/tests/sentry/api/bases/test_project.py b/tests/sentry/api/bases/test_project.py index 582d62a7f3d596..1fb2aa27f8aa19 100644 --- a/tests/sentry/api/bases/test_project.py +++ b/tests/sentry/api/bases/test_project.py @@ -1,6 +1,9 @@ +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 @@ -8,6 +11,7 @@ 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) diff --git a/tests/sentry/api/bases/test_team.py b/tests/sentry/api/bases/test_team.py index 822a8f5d0f6c3d..771d54a47911d4 100644 --- a/tests/sentry/api/bases/test_team.py +++ b/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() diff --git a/tests/sentry/issues/endpoints/test_group.py b/tests/sentry/issues/endpoints/test_group.py index 1cf4eae09c39cd..9bc229883ac4f1 100644 --- a/tests/sentry/issues/endpoints/test_group.py +++ b/tests/sentry/issues/endpoints/test_group.py @@ -2,6 +2,7 @@ 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 @@ -9,6 +10,7 @@ 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 diff --git a/tests/sentry/test_viewer_context.py b/tests/sentry/test_viewer_context.py index 9ac2458c7dd921..e1c3ecb56dbf4e 100644 --- a/tests/sentry/test_viewer_context.py +++ b/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