ref(mediators): turn GrantExchanger into a dataclass (#78952)

src/sentry/mediators/__init__.py

@@ -1,5 +1,4 @@
 from .mediator import Mediator  # NOQA
 from .param import Param  # NOQA
-from .token_exchange.grant_exchanger import GrantExchanger  # noqa: F401
 from .token_exchange.refresher import Refresher  # noqa: F401
 from .token_exchange.util import AUTHORIZATION, REFRESH, GrantTypes  # noqa: F401

src/sentry/mediators/token_exchange/__init__.py

@@ -1,4 +1,3 @@
-from .grant_exchanger import GrantExchanger  # NOQA
 from .refresher import Refresher  # NOQA
 from .util import AUTHORIZATION, REFRESH, GrantTypes, token_expiration  # NOQA
 from .validator import Validator  # NOQA

src/sentry/sentry_apps/api/endpoints/sentry_app_authorizations.py

@@ -10,10 +10,10 @@
 from sentry.api.serializers.models.apitoken import ApiTokenSerializer
 from sentry.auth.services.auth.impl import promote_request_api_user
 from sentry.coreapi import APIUnauthorized
-from sentry.mediators.token_exchange.grant_exchanger import GrantExchanger
 from sentry.mediators.token_exchange.refresher import Refresher
 from sentry.mediators.token_exchange.util import GrantTypes
 from sentry.sentry_apps.api.bases.sentryapps import SentryAppAuthorizationsBaseEndpoint
+from sentry.sentry_apps.token_exchange.grant_exchanger import GrantExchanger
 
 logger = logging.getLogger(__name__)
 
@@ -34,12 +34,12 @@ def post(self, request: Request, installation) -> Response:
 
         try:
             if request.json_body.get("grant_type") == GrantTypes.AUTHORIZATION:
-                token = GrantExchanger.run(
+                token = GrantExchanger(
                     install=installation,
                     code=request.json_body.get("code"),
                     client_id=request.json_body.get("client_id"),
                     user=promote_request_api_user(request),
-                )
+                ).run()
             elif request.json_body.get("grant_type") == GrantTypes.REFRESH:
                 token = Refresher.run(
                     install=installation,

src/sentry/mediators/token_exchange/grant_exchanger.py → src/sentry/sentry_apps/token_exchange/grant_exchanger.py

@@ -1,12 +1,11 @@
+from dataclasses import dataclass
 from datetime import datetime, timezone
 
-from django.db import router
+from django.db import router, transaction
 from django.utils.functional import cached_property
 
 from sentry import analytics
 from sentry.coreapi import APIUnauthorized
-from sentry.mediators.mediator import Mediator
-from sentry.mediators.param import Param
 from sentry.mediators.token_exchange.util import token_expiration
 from sentry.mediators.token_exchange.validator import Validator
 from sentry.models.apiapplication import ApiApplication
@@ -19,71 +18,75 @@
 from sentry.users.models.user import User
 
 
-class GrantExchanger(Mediator):
+@dataclass
+class GrantExchanger:
     """
     Exchanges a Grant Code for an Access Token
     """
 
-    install = Param(RpcSentryAppInstallation)
-    code = Param(str)
-    client_id = Param(str)
-    user = Param(User)
-    using = router.db_for_write(User)
+    install: RpcSentryAppInstallation
+    code: str
+    client_id: str
+    user: User
 
-    def call(self):
-        self._validate()
-        self._create_token()
+    def run(self):
+        with transaction.atomic(using=router.db_for_write(ApiToken)):
+            self._validate()
+            token = self._create_token()
 
-        # Once it's exchanged it's no longer valid and should not be
-        # exchangeable, so we delete it.
-        self._delete_grant()
+            # Once it's exchanged it's no longer valid and should not be
+            # exchangeable, so we delete it.
+            self._delete_grant()
+            self.record_analytics()
 
-        return self.token
+            return token
 
-    def record_analytics(self):
+    def record_analytics(self) -> None:
         analytics.record(
             "sentry_app.token_exchanged",
             sentry_app_installation_id=self.install.id,
             exchange_type="authorization",
         )
 
-    def _validate(self):
+    def _validate(self) -> None:
         Validator.run(install=self.install, client_id=self.client_id, user=self.user)
 
         if not self._grant_belongs_to_install() or not self._sentry_app_user_owns_grant():
-            raise APIUnauthorized
+            raise APIUnauthorized("Forbidden grant")
 
         if not self._grant_is_active():
             raise APIUnauthorized("Grant has already expired.")
 
-    def _grant_belongs_to_install(self):
+    def _grant_belongs_to_install(self) -> bool:
         return self.grant.sentry_app_installation.id == self.install.id
 
-    def _sentry_app_user_owns_grant(self):
+    def _sentry_app_user_owns_grant(self) -> bool:
         return self.grant.application.owner == self.user
 
-    def _grant_is_active(self):
+    def _grant_is_active(self) -> bool:
         return self.grant.expires_at > datetime.now(timezone.utc)
 
-    def _delete_grant(self):
+    def _delete_grant(self) -> None:
         # This will cause a set null to trigger which does not need to cascade an outbox
         with unguarded_write(router.db_for_write(ApiGrant)):
             self.grant.delete()
 
-    def _create_token(self):
-        self.token = ApiToken.objects.create(
+    def _create_token(self) -> ApiToken:
+        token = ApiToken.objects.create(
             user=self.user,
             application=self.application,
             scope_list=self.sentry_app.scope_list,
             expires_at=token_expiration(),
         )
         try:
-            SentryAppInstallation.objects.get(id=self.install.id).update(api_token=self.token)
+            SentryAppInstallation.objects.get(id=self.install.id).update(api_token=token)
         except SentryAppInstallation.DoesNotExist:
             pass
 
+        return token
+
     @cached_property
-    def grant(self):
+    def grant(self) -> ApiGrant:
         try:
             return (
                 ApiGrant.objects.select_related("sentry_app_installation")
@@ -92,18 +95,18 @@ def grant(self):
                 .get(code=self.code)
             )
         except ApiGrant.DoesNotExist:
-            raise APIUnauthorized
+            raise APIUnauthorized("Could not find grant")
 
     @property
-    def application(self):
+    def application(self) -> ApiApplication:
         try:
             return self.grant.application
         except ApiApplication.DoesNotExist:
-            raise APIUnauthorized
+            raise APIUnauthorized("Could not find application")
 
     @property
-    def sentry_app(self):
+    def sentry_app(self) -> SentryApp:
         try:
             return self.application.sentry_app
         except SentryApp.DoesNotExist:
-            raise APIUnauthorized
+            raise APIUnauthorized("Could not find sentry app")

src/sentry/testutils/factories.py

@@ -75,7 +75,6 @@
 from sentry.integrations.models.repository_project_path_config import RepositoryProjectPathConfig
 from sentry.integrations.types import ExternalProviders
 from sentry.issues.grouptype import get_group_type_by_type_id
-from sentry.mediators.token_exchange.grant_exchanger import GrantExchanger
 from sentry.models.activity import Activity
 from sentry.models.apikey import ApiKey
 from sentry.models.apitoken import ApiToken
@@ -144,6 +143,7 @@
 from sentry.sentry_apps.models.servicehook import ServiceHook
 from sentry.sentry_apps.services.app.serial import serialize_sentry_app_installation
 from sentry.sentry_apps.services.hook import hook_service
+from sentry.sentry_apps.token_exchange.grant_exchanger import GrantExchanger
 from sentry.signals import project_created
 from sentry.silo.base import SiloMode
 from sentry.snuba.dataset import Dataset
@@ -1214,12 +1214,13 @@ def create_sentry_app_installation(
             ):
                 assert install.api_grant is not None
                 assert install.sentry_app.application is not None
-                GrantExchanger.run(
+                assert install.sentry_app.proxy_user is not None
+                GrantExchanger(
                     install=rpc_install,
                     code=install.api_grant.code,
                     client_id=install.sentry_app.application.client_id,
                     user=install.sentry_app.proxy_user,
-                )
+                ).run()
                 install = SentryAppInstallation.objects.get(id=install.id)
         return install
 

tests/sentry/sentry_apps/api/endpoints/test_organization_sentry_app_installation_details.py

@@ -5,9 +5,9 @@
 
 from sentry import audit_log
 from sentry.constants import SentryAppInstallationStatus
-from sentry.mediators.token_exchange.grant_exchanger import GrantExchanger
 from sentry.models.auditlogentry import AuditLogEntry
 from sentry.sentry_apps.services.app import app_service
+from sentry.sentry_apps.token_exchange.grant_exchanger import GrantExchanger
 from sentry.testutils.cases import APITestCase
 from sentry.testutils.silo import control_silo_test
 from sentry.users.services.user.service import user_service
@@ -141,12 +141,12 @@ def test_member_cannot_delete_install(self):
 class MarkInstalledSentryAppInstallationsTest(SentryAppInstallationDetailsTest):
     def setUp(self):
         super().setUp()
-        self.token = GrantExchanger.run(
+        self.token = GrantExchanger(
             install=self.installation,
             code=self.orm_installation.api_grant.code,
             client_id=self.published_app.application.client_id,
             user=self.published_app.proxy_user,
-        )
+        ).run()
 
     @patch("sentry.analytics.record")
     def test_sentry_app_installation_mark_installed(self, record):
@@ -188,12 +188,12 @@ def test_sentry_app_installation_mark_pending_status(self):
         )
 
     def test_sentry_app_installation_mark_installed_wrong_app(self):
-        self.token = GrantExchanger.run(
+        self.token = GrantExchanger(
             install=self.installation2,
             code=self.orm_installation2.api_grant.code,
             client_id=self.unpublished_app.application.client_id,
             user=self.unpublished_app.proxy_user,
-        )
+        ).run()
         self.url = reverse(
             "sentry-api-0-sentry-app-installation-details", args=[self.installation.uuid]
         )

tests/sentry/mediators/token_exchange/test_grant_exchanger.py → tests/sentry/sentry_apps/token_exchange/test_grant_exchanger.py

@@ -4,12 +4,12 @@
 import pytest
 
 from sentry.coreapi import APIUnauthorized
-from sentry.mediators.token_exchange.grant_exchanger import GrantExchanger
 from sentry.models.apiapplication import ApiApplication
 from sentry.models.apigrant import ApiGrant
 from sentry.sentry_apps.models.sentry_app import SentryApp
 from sentry.sentry_apps.models.sentry_app_installation import SentryAppInstallation
 from sentry.sentry_apps.services.app import app_service
+from sentry.sentry_apps.token_exchange.grant_exchanger import GrantExchanger
 from sentry.testutils.cases import TestCase
 from sentry.testutils.silo import control_silo_test
 
@@ -29,15 +29,15 @@ def setUp(self):
         )
 
     def test_happy_path(self):
-        assert self.grant_exchanger.call()
+        assert self.grant_exchanger.run()
 
     def test_adds_token_to_installation(self):
-        token = self.grant_exchanger.call()
+        token = self.grant_exchanger.run()
         assert SentryAppInstallation.objects.get(id=self.install.id).api_token == token
 
     @patch("sentry.mediators.token_exchange.Validator.run")
     def test_validate_generic_token_exchange_requirements(self, validator):
-        self.grant_exchanger.call()
+        self.grant_exchanger.run()
 
         validator.assert_called_once_with(
             install=self.install, client_id=self.client_id, user=self.user
@@ -48,46 +48,46 @@ def test_grant_must_belong_to_installations(self):
         self.grant_exchanger.code = other_install.api_grant.code
 
         with pytest.raises(APIUnauthorized):
-            self.grant_exchanger.call()
+            self.grant_exchanger.run()
 
     def test_request_user_owns_api_grant(self):
         self.grant_exchanger.user = self.create_user()
 
         with pytest.raises(APIUnauthorized):
-            self.grant_exchanger.call()
+            self.grant_exchanger.run()
 
     def test_grant_must_be_active(self):
         self.orm_install.api_grant.update(expires_at=(datetime.now(UTC) - timedelta(hours=1)))
 
         with pytest.raises(APIUnauthorized):
-            self.grant_exchanger.call()
+            self.grant_exchanger.run()
 
     def test_grant_must_exist(self):
         self.grant_exchanger.code = "123"
 
         with pytest.raises(APIUnauthorized):
-            self.grant_exchanger.call()
+            self.grant_exchanger.run()
 
     @patch("sentry.models.ApiGrant.application", side_effect=ApiApplication.DoesNotExist)
     def test_application_must_exist(self, _):
         with pytest.raises(APIUnauthorized):
-            self.grant_exchanger.call()
+            self.grant_exchanger.run()
 
     @patch("sentry.models.ApiApplication.sentry_app", side_effect=SentryApp.DoesNotExist)
     def test_sentry_app_must_exist(self, _):
         with pytest.raises(APIUnauthorized):
-            self.grant_exchanger.call()
+            self.grant_exchanger.run()
 
     def test_deletes_grant_on_successful_exchange(self):
         grant_id = self.orm_install.api_grant_id
-        self.grant_exchanger.call()
+        self.grant_exchanger.run()
         assert not ApiGrant.objects.filter(id=grant_id)
 
     @patch("sentry.analytics.record")
     def test_records_analytics(self, record):
-        GrantExchanger.run(
+        GrantExchanger(
             install=self.install, client_id=self.client_id, code=self.code, user=self.user
-        )
+        ).run()
 
         record.assert_called_with(
             "sentry_app.token_exchanged",