Cache the current context for the thread

Use a threading.local instance to store the current RequestContext, with
an option to not overwrite an existing context.

bp/graduate-oslo-context

Change-Id: I000cb13392ee21258dc2a91683294dc9ff2aeb8f

2 files changed, 62 insertions, 2 deletions

diff --git a/oslo_context/context.py b/oslo_context/context.py
index 58425d3..44b2b36 100644
--- a/oslo_context/context.py
+++ b/oslo_context/context.py
@@ -21,9 +21,13 @@ context or provide additional information in their specific WSGI pipeline.
 """
 
 import itertools
+import threading
 import uuid
 
 
+_request_store = threading.local()
+
+
 def generate_request_id():
     return b'req-' + str(uuid.uuid4()).encode('ascii')
 
@@ -41,7 +45,12 @@ class RequestContext(object):
     def __init__(self, auth_token=None, user=None, tenant=None, domain=None,
                  user_domain=None, project_domain=None, is_admin=False,
                  read_only=False, show_deleted=False, request_id=None,
-                 resource_uuid=None):
+                 resource_uuid=None, overwrite=True):
+        """Initialize the RequestContext
+
+        :param overwrite: Set to False to ensure that the greenthread local
+                          copy of the index is not overwritten.
+        """
         self.auth_token = auth_token
         self.user = user
         self.tenant = tenant
@@ -55,6 +64,11 @@ class RequestContext(object):
         if not request_id:
             request_id = generate_request_id()
         self.request_id = request_id
+        if overwrite or not get_current():
+            self.update_store()
+
+    def update_store(self):
+        _request_store.context = self
 
     def to_dict(self):
         user_idt = (
@@ -98,7 +112,8 @@ def get_admin_context(show_deleted=False):
     context = RequestContext(None,
                              tenant=None,
                              is_admin=True,
-                             show_deleted=show_deleted)
+                             show_deleted=show_deleted,
+                             overwrite=False)
     return context
 
 
@@ -125,3 +140,11 @@ def is_user_context(context):
     if not context.user_id or not context.project_id:
         return False
     return True
+
+
+def get_current():
+    """Return this thread's current context
+
+    If no context is set, returns None
+    """
+    return getattr(_request_store, 'context', None)
diff --git a/oslo_context/tests/test_context.py b/oslo_context/tests/test_context.py
index 158ba8d..655895d 100644
--- a/oslo_context/tests/test_context.py
+++ b/oslo_context/tests/test_context.py
@@ -20,10 +20,47 @@ from oslo_context import context
 
 class ContextTest(test_base.BaseTestCase):
 
+    def setUp(self):
+        super(ContextTest, self).setUp()
+        self.addCleanup(self._remove_cached_context)
+
+    def _remove_cached_context(self):
+        """Remove the thread-local context stored in the module."""
+        try:
+            del context._request_store.context
+        except AttributeError:
+            pass
+
     def test_context(self):
         ctx = context.RequestContext()
         self.assertTrue(ctx)
 
+    def test_store_when_no_overwrite(self):
+        # If no context exists we store one even if overwrite is false
+        # (since we are not overwriting anything).
+        ctx = context.RequestContext(overwrite=False)
+        self.assertIs(context.get_current(), ctx)
+
+    def test_no_overwrite(self):
+        # If there is already a context in the cache a new one will
+        # not overwrite it if overwrite=False.
+        ctx1 = context.RequestContext(overwrite=True)
+        context.RequestContext(overwrite=False)
+        self.assertIs(context.get_current(), ctx1)
+
+    def test_admin_no_overwrite(self):
+        # If there is already a context in the cache creating an admin
+        # context will not overwrite it.
+        ctx1 = context.RequestContext(overwrite=True)
+        context.get_admin_context()
+        self.assertIs(context.get_current(), ctx1)
+
+    def test_store_current(self):
+        # By default a new context is stored.
+        self._remove_cached_context()
+        ctx = context.RequestContext()
+        self.assertIs(context.get_current(), ctx)
+
     def test_admin_context_show_deleted_flag_default(self):
         ctx = context.get_admin_context()
         self.assertFalse(ctx.show_deleted)