Refactor context module to make it more flexible

- context data is stored in a dictionary now, so any key is allowed;
- list of allowed key for certain context is stored in special class
  variable;
- methods "current" and "elevated" of ClimateContext class are
  accessible from module level;
- moved to usage of threading.local instead of emulating one in code (as
  long as we use eventlet's monkey-patching, it's specific local() will
  be used).

Change-Id: I3c79c9cd0d16dfb5c61208f2aa71b5c59f3f40f3

climate/api/app.py

@@ -21,7 +21,6 @@ from werkzeug import exceptions as werkzeug_exceptions
 
 from climate.api import utils as api_utils
 from climate.api import v1_0 as api_v1_0
-from climate import context
 from climate.openstack.common import log
 from climate.openstack.common.middleware import debug
 
@@ -61,10 +60,6 @@ def version_list():
     })
 
 
-def teardown_request(_ex=None):
-    context.Context.clear()
-
-
 def make_app():
     """App builder (wsgi).
 
@@ -73,7 +68,6 @@ def make_app():
     app = flask.Flask('climate.api')
 
     app.route('/', methods=['GET'])(version_list)
-    app.teardown_request(teardown_request)
     app.register_blueprint(api_v1_0.rest, url_prefix='/v1')
 
     for code in werkzeug_exceptions.default_exceptions.iterkeys():

climate/api/context.py

@@ -17,7 +17,7 @@ from climate import context
 
 
 def ctx_from_headers(headers):
-    return context.Context(
+    return context.ClimateContext(
         user_id=headers['X-User-Id'],
         tenant_id=headers['X-Tenant-Id'],
         auth_token=headers['X-Auth-Token'],

climate/context.py

@@ -13,73 +13,86 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from eventlet import corolocal
+import threading
-
-from climate.openstack.common import log as logging
 
 
-LOG = logging.getLogger(__name__)
+class BaseContext(object):
 
+    _elements = set()
+    _context_stack = threading.local()
 
-class Context(object):
+    def __init__(self, __mapping=None, **kwargs):
-    """Context class for the Climate operations."""
+        if __mapping is None:
-    _contexts = {}
+            self.__values = dict(**kwargs)
+        else:
+            if isinstance(__mapping, BaseContext):
+                __mapping = __mapping.__values
+            self.__values = dict(__mapping)
+            self.__values.update(**kwargs)
+        bad_keys = set(self.__values) - self._elements
+        if bad_keys:
+            raise TypeError("Only %s keys are supported. %s given" %
+                            (tuple(self._elements), tuple(bad_keys)))
 
-    def __init__(self, user_id=None, tenant_id=None, auth_token=None,
+    def __getattr__(self, name):
-                 service_catalog=None, user_name=None, tenant_name=None,
+        try:
-                 roles=None, **kwargs):
+            return self.__values[name]
-        if kwargs:
+        except KeyError:
-            LOG.warn('Arguments dropped when creating context: %s', kwargs)
+            if name in self._elements:
-
+                return None
-        self.user_id = user_id
+            else:
-        self.user_name = user_name
+                raise AttributeError(name)
-        self.tenant_id = tenant_id
-        self.tenant_name = tenant_name
-        self.auth_token = auth_token
-        self.service_catalog = service_catalog
-        self.roles = roles
-        self._db_session = None
 
     def __enter__(self):
-        stack = self._contexts.setdefault(corolocal.get_ident(), [])
+        try:
+            stack = self._context_stack.stack
+        except AttributeError:
+            stack = []
+            self._context_stack.stack = stack
         stack.append(self)
+        return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        stack = self._contexts[corolocal.get_ident()]
+        res = self._context_stack.stack.pop()
-        stack.pop()
+        assert res is self, "self should be the top element of the stack"
-        if not stack:
-            del self._contexts[corolocal.get_ident()]
 
     @classmethod
     def current(cls):
         try:
-            return cls._contexts[corolocal.get_ident()][-1]
+            return cls._context_stack.stack[-1]
-        except (KeyError, IndexError):
+        except (AttributeError, IndexError):
             raise RuntimeError("Context isn't available here")
 
-    @classmethod
+    # NOTE(yorik-sar): as long as oslo.rpc requires this
-    def clear(cls):
-        try:
-            del cls._contexts[corolocal.get_ident()]
-        except KeyError:
-            pass
-
-    def clone(self):
-        return Context(self.user_id,
-                       self.tenant_id,
-                       self.auth_token,
-                       self.service_catalog,
-                       self.user_name,
-                       self.tenant_name,
-                       self.roles)
-
     def to_dict(self):
-        return {
+        return self.__values
-            'user_id': self.user_id,
+
-            'user_name': self.user_name,
+
-            'tenant_id': self.tenant_id,
+class ClimateContext(BaseContext):
-            'tenant_name': self.tenant_name,
+
-            'auth_token': self.auth_token,
+    _elements = set([
-            'service_catalog': self.service_catalog,
+        "user_id",
-            'roles': self.roles,
+        "tenant_id",
-        }
+        "auth_token",
+        "service_catalog",
+        "user_name",
+        "tenant_name",
+        "roles",
+        "is_admin",
+    ])
+
+    @classmethod
+    def elevated(cls):
+        try:
+            ctx = cls.current()
+        except RuntimeError:
+            ctx = None
+        return cls(ctx, is_admin=True)
+
+
+def current():
+    return ClimateContext.current()
+
+
+def elevated():
+    return ClimateContext.elevated()

climate/db/sqlalchemy/api.py

@@ -51,7 +51,7 @@ def model_query(model, session=None, project_only=None):
     query = session.query(model)
 
     if project_only:
-        ctx = context.Context.current()
+        ctx = context.current()
         query = query.filter_by(tenant_id=ctx.project_id)
 
     return query
@@ -63,7 +63,7 @@ def column_query(*columns, **kwargs):
     query = session.query(*columns)
 
     if kwargs.get("project_only"):
-        ctx = context.Context.current()
+        ctx = context.current()
         query = query.filter_by(tenant_id=ctx.tenant_id)
 
     return query

climate/tests/__init__.py

@@ -74,7 +74,7 @@ class TestCase(test.BaseTestCase):
 
     def set_context(self, ctx):
         if self.context_mock is None:
-            self.context_mock = self.patch(context.Context, 'current')
+            self.context_mock = self.patch(context.ClimateContext, 'current')
         self.context_mock.return_value = ctx
 
 

climate/tests/api/test_app.py

@@ -19,7 +19,6 @@ from werkzeug import exceptions as werkzeug_exceptions
 
 from climate.api import app
 from climate.api import utils as api_utils
-from climate import context
 from climate import tests
 
 
@@ -29,12 +28,10 @@ class AppTestCase(tests.TestCase):
 
         self.app = app
         self.api_utils = api_utils
-        self.context = context
         self.flask = flask
         self.auth_token = auth_token
 
         self.render = self.patch(self.api_utils, 'render')
-        self.context_clear = self.patch(self.context.Context, 'clear')
         self.fake_app = self.patch(self.flask, 'Flask')
         self.fake_ff = self.patch(self.auth_token, 'filter_factory')
 
@@ -62,10 +59,6 @@ class AppTestCase(tests.TestCase):
             ],
         })
 
-    def test_teardown_request(self):
-        self.app.teardown_request()
-        self.context_clear.assert_called_once()
-
     def test_make_app(self):
         self.app.make_app()
         self.fake_ff.assert_called_once_with(self.fake_app().config,

climate/tests/api/test_context.py

@@ -22,7 +22,7 @@ class ContextTestCase(tests.TestCase):
     def setUp(self):
         super(ContextTestCase, self).setUp()
 
-        self.context = self.patch(context, 'Context')
+        self.context = self.patch(context, 'ClimateContext')
         self.fake_headers = {u'X-User-Id': u'1',
                              u'X-Tenant-Id': u'1',
                              u'X-Auth-Token': u'111-111-111',

climate/tests/test_context.py

@@ -0,0 +1,104 @@
+# Copyright (c) 2013 Mirantis Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from climate import context
+from climate import tests
+
+
+class TestContext(context.BaseContext):
+    _elements = set(["first", "second", "third"])
+
+
+class TestContextCreate(tests.TestCase):
+
+    def test_kwargs(self):
+        ctx = TestContext(first=1, second=2)
+        self.assertEqual(ctx.to_dict(), {"first": 1, "second": 2})
+
+    def test_dict(self):
+        ctx = TestContext({"first": 1, "second": 2})
+        self.assertEqual(ctx.to_dict(), {"first": 1, "second": 2})
+
+    def test_mix(self):
+        ctx = TestContext({"first": 1}, second=2)
+        self.assertEqual(ctx.to_dict(), {"first": 1, "second": 2})
+
+    def test_fail(self):
+        self.assertRaises(TypeError, TestContext, forth=4)
+
+
+class TestBaseContext(tests.TestCase):
+
+    def setUp(self):
+        super(TestBaseContext, self).setUp()
+        self.context = TestContext(first=1, second=2)
+
+    def test_get_defined(self):
+        super(TestBaseContext, self).tearDown()
+        self.assertEqual(self.context.first, 1)
+
+    def test_get_default(self):
+        self.assertIsNone(self.context.third)
+
+    def test_get_unexpected(self):
+        self.assertRaises(AttributeError, getattr, self.context, 'forth')
+
+    def test_current_fails(self):
+        self.assertRaises(RuntimeError, TestContext.current)
+
+
+class TestContextManager(tests.TestCase):
+
+    def setUp(self):
+        super(TestContextManager, self).setUp()
+        self.context = TestContext(first=1, second=2)
+        self.context.__enter__()
+
+    def tearDown(self):
+        super(TestContextManager, self).tearDown()
+        self.context.__exit__(None, None, None)
+        try:
+            stack = TestContext._context_stack.stack
+        except AttributeError:
+            self.fail("Context stack have never been created")
+        else:
+            del TestContext._context_stack.stack
+            self.assertEqual(stack, [],
+                             "Context stack is not empty after test.")
+
+    def test_enter(self):
+        self.assertEqual(TestContext._context_stack.stack, [self.context])
+
+    def test_double_enter(self):
+        with self.context:
+            self.assertEqual(TestContext._context_stack.stack,
+                             [self.context, self.context])
+
+    def test_current(self):
+        self.assertIs(self.context, TestContext.current())
+
+
+class TestClimateContext(tests.TestCase):
+
+    def test_elevated_empty(self):
+        ctx = context.ClimateContext.elevated()
+        self.assertEqual(ctx.is_admin, True)
+
+    def test_elevated(self):
+        with context.ClimateContext(user_id="user", tenant_id="tenant"):
+            ctx = context.ClimateContext.elevated()
+            self.assertEqual(ctx.user_id, "user")
+            self.assertEqual(ctx.tenant_id, "tenant")
+            self.assertEqual(ctx.is_admin, True)

climate/utils/service.py

@@ -28,14 +28,14 @@ import climate.openstack.common.rpc.proxy as rpc_proxy
 class RpcProxy(rpc_proxy.RpcProxy):
     def cast(self, name, topic=None, version=None, ctx=None, **kwargs):
         if ctx is None:
-            ctx = context.Context.current()
+            ctx = context.current()
         msg = self.make_msg(name, **kwargs)
         return super(RpcProxy, self).cast(ctx, msg,
                                           topic=topic, version=version)
 
     def call(self, name, topic=None, version=None, ctx=None, **kwargs):
         if ctx is None:
-            ctx = context.Context.current()
+            ctx = context.current()
         msg = self.make_msg(name, **kwargs)
         return super(RpcProxy, self).call(ctx, msg,
                                           topic=topic, version=version)
@@ -45,9 +45,9 @@ def export_context(func):
     @functools.wraps(func)
     def decorator(manager, ctx, *args, **kwargs):
         try:
-            context.Context.current()
+            context.current()
         except RuntimeError:
-            new_ctx = context.Context(**ctx.values)
+            new_ctx = context.ClimateContext(ctx.values)
             with new_ctx:
                 return func(manager, *args, **kwargs)
         else:
@@ -59,7 +59,7 @@ def export_context(func):
 def with_empty_context(func):
     @functools.wraps(func)
     def decorator(*args, **kwargs):
-        with context.Context():
+        with context.ClimateContext():
             return func(*args, **kwargs)
 
     return decorator