feat(auth): adds request_cache decorator to cache org member (#26165)

Author: Stephen Cefali

src/sentry/auth/access.py

@@ -19,6 +19,12 @@
     Team,
     UserPermission,
 )
+from sentry.utils.request_cache import request_cache
+
+
+@request_cache
+def get_cached_organization_member(user_id, organization_id):
+    return OrganizationMember.objects.get(user_id=user_id, organization_id=organization_id)
 
 
 def _sso_params(member):
@@ -279,7 +285,7 @@ def from_request(request, organization=None, scopes=None):
         # we special case superuser so that if they're a member of the org
         # they must still follow SSO checks, but they gain global access
         try:
-            member = OrganizationMember.objects.get(user=request.user, organization=organization)
+            member = get_cached_organization_member(request.user.id, organization.id)
         except OrganizationMember.DoesNotExist:
             requires_sso, sso_is_valid = False, True
         else:
@@ -345,7 +351,7 @@ def from_user(user, organization=None, scopes=None):
         return OrganizationlessAccess(permissions=UserPermission.for_user(user.id))
 
     try:
-        om = OrganizationMember.objects.get(user=user, organization=organization)
+        om = get_cached_organization_member(user.id, organization.id)
     except OrganizationMember.DoesNotExist:
         return OrganizationlessAccess(permissions=UserPermission.for_user(user.id))
 

src/sentry/utils/request_cache.py

@@ -0,0 +1,42 @@
+import threading
+
+from celery.signals import task_failure, task_success
+from django.core.signals import request_finished
+
+from sentry import app
+
+_cache = threading.local()
+
+
+def request_cache(func):
+    """
+    A decorator to memoize functions on a per-request basis.
+    Arguments to the memoized function should NOT be objects
+    Use primitive types as arguments
+    """
+
+    def wrapped(*args, **kwargs):
+        # if no request, skip cache
+        if app.env.request is None:
+            return func(*args, **kwargs)
+
+        if not hasattr(_cache, "items"):
+            _cache.items = {}
+        cache_key = (func, repr(args), repr(kwargs))
+        if cache_key in _cache.items:
+            rv = _cache.items[cache_key]
+        else:
+            rv = func(*args, **kwargs)
+            _cache.items[cache_key] = rv
+        return rv
+
+    return wrapped
+
+
+def clear_cache(**kwargs):
+    _cache.items = {}
+
+
+request_finished.connect(clear_cache)
+task_failure.connect(clear_cache)
+task_success.connect(clear_cache)

tests/sentry/utils/test_request_cache.py

@@ -0,0 +1,78 @@
+from django.core.handlers.wsgi import WSGIHandler
+from django.core.signals import request_finished
+from django.http import HttpRequest
+from django.utils import timezone
+
+from sentry import app
+from sentry.testutils import TestCase
+from sentry.utils.compat.mock import patch
+from sentry.utils.request_cache import clear_cache, request_cache
+
+
+@request_cache
+def cached_fn(arg1, arg2=None):
+    # call this so we can mock it and see how often it's called
+    timezone.now()
+    if arg2:
+        return arg1 + arg2
+    return arg1
+
+
+class RequestCacheTest(TestCase):
+    def setUp(self):
+        self.original_receivers = request_finished.receivers
+        request_finished.receivers = []
+        request_finished.connect(clear_cache)
+        super().setUp()
+
+    def tearDown(self):
+        # remove the request, trigger the signal to clear the cache, restore receivers
+        app.env.request = None
+        request_finished.send(sender=WSGIHandler)
+        request_finished.receivers = self.original_receivers
+        super().tearDown()
+
+    @patch("django.utils.timezone.now")
+    def test_basic_cache(self, mock_now):
+        app.env.request = HttpRequest()
+        assert cached_fn("cat", arg2="dog") == "catdog"
+        assert cached_fn("cat", arg2="dog") == "catdog"
+        assert mock_now.call_count == 1
+
+    @patch("django.utils.timezone.now")
+    def test_cache_none(self, mock_now):
+        app.env.request = HttpRequest()
+        assert cached_fn(None) is None
+        assert cached_fn(None) is None
+        assert mock_now.call_count == 1
+
+    @patch("django.utils.timezone.now")
+    def test_different_args(self, mock_now):
+        app.env.request = HttpRequest()
+        assert cached_fn("cat", arg2="dog") == "catdog"
+        assert cached_fn("hey", arg2="dog") == "heydog"
+        assert mock_now.call_count == 2
+
+    @patch("django.utils.timezone.now")
+    def test_different_kwargs(self, mock_now):
+        app.env.request = HttpRequest()
+        assert cached_fn("cat", arg2="dog") == "catdog"
+        assert cached_fn("cat", arg2="hat") == "cathat"
+        assert mock_now.call_count == 2
+
+    @patch("django.utils.timezone.now")
+    def test_different_request(self, mock_now):
+        app.env.request = HttpRequest()
+        assert cached_fn("cat") == "cat"
+        request_finished.send(sender=WSGIHandler)
+        app.env.request = HttpRequest()
+        assert cached_fn("cat") == "cat"
+        assert mock_now.call_count == 2
+
+    @patch("django.utils.timezone.now")
+    def test_request_over(self, mock_now):
+        app.env.request = HttpRequest()
+        assert cached_fn("cat") == "cat"
+        app.env.request = None
+        assert cached_fn("cat") == "cat"
+        assert mock_now.call_count == 2