Added allow_async property

In Nova we need a decorator that is READER by default
but can be used as ASYNC if asynchronous session is started before.

Co-Authored-By: Sergey Nikitin <snikitin@mirantis.com>

Change-Id: I0c0f332d87bf6b1213c800763918340b1e226934

2 files changed, 84 insertions, 6 deletions

diff --git a/oslo_db/sqlalchemy/enginefacade.py b/oslo_db/sqlalchemy/enginefacade.py
index 7f64acc..2b3fa73 100644
--- a/oslo_db/sqlalchemy/enginefacade.py
+++ b/oslo_db/sqlalchemy/enginefacade.py
@@ -528,13 +528,13 @@ class _TransactionContext(object):
         else:
             transaction.rollback()
 
-    def _produce_block(self, mode, connection, savepoint):
+    def _produce_block(self, mode, connection, savepoint, allow_async=False):
         if mode is _WRITER:
             self._writer()
         elif mode is _ASYNC_READER:
             self._async_reader()
         else:
-            self._reader()
+            self._reader(allow_async)
         if connection:
             return self._connection(savepoint)
         else:
@@ -552,10 +552,10 @@ class _TransactionContext(object):
                 "Can't upgrade an ASYNC_READER transaction "
                 "to a WRITER mid-transaction")
 
-    def _reader(self):
+    def _reader(self, allow_async=False):
         if self.mode is None:
             self.mode = _READER
-        elif self.mode is _ASYNC_READER:
+        elif self.mode is _ASYNC_READER and not allow_async:
             raise TypeError(
                 "Can't upgrade an ASYNC_READER transaction "
                 "to a READER mid-transaction")
@@ -589,7 +589,8 @@ class _TransactionContextManager(object):
             savepoint=False,
             connection=False,
             replace_global_factory=None,
-            _is_global_manager=False):
+            _is_global_manager=False,
+            allow_async=False):
 
         if root is None:
             self._root = self
@@ -606,6 +607,7 @@ class _TransactionContextManager(object):
             raise TypeError(
                 "setting savepoint and independent makes no sense.")
         self._connection = connection
+        self._allow_async = allow_async
 
     @property
     def _factory(self):
@@ -649,6 +651,27 @@ class _TransactionContextManager(object):
         return self._clone(mode=_READER)
 
     @property
+    def allow_async(self):
+        """Modifier to allow async operations
+
+        Allows async operations if asynchronous session is already
+        started in this context. Marking DB API methods with READER would make
+        it impossible to use them in ASYNC_READER transactions, and marking
+        them with ASYNC_READER would require a modification of all the places
+        these DB API methods are called to force READER mode, where the latest
+        DB state is required.
+
+        In Nova DB API methods should have a 'safe' default (i.e. READER),
+        so that they can start sessions on their own, but it would also be
+        useful for them to be able to participate in an existing ASYNC_READER
+        session, if one was started up the stack.
+        """
+
+        if self._mode is _WRITER:
+            raise TypeError("Setting async on a WRITER makes no sense")
+        return self._clone(allow_async=True)
+
+    @property
     def independent(self):
         """Modifier to start a transaction independent from any enclosing."""
         return self._clone(independent=True)
@@ -732,7 +755,8 @@ class _TransactionContextManager(object):
                 with current._produce_block(
                     mode=self._mode,
                     connection=self._connection,
-                    savepoint=self._savepoint) as resource:
+                    savepoint=self._savepoint,
+                    allow_async=self._allow_async) as resource:
                     yield resource
             else:
                 yield
diff --git a/oslo_db/tests/sqlalchemy/test_enginefacade.py b/oslo_db/tests/sqlalchemy/test_enginefacade.py
index 9226c05..97229f0 100644
--- a/oslo_db/tests/sqlalchemy/test_enginefacade.py
+++ b/oslo_db/tests/sqlalchemy/test_enginefacade.py
@@ -583,6 +583,60 @@ class MockFacadeTest(oslo_test_base.BaseTestCase):
             exc.args[0]
         )
 
+    def test_reader_allow_async_nested_in_async_reader(self):
+        context = oslo_context.RequestContext()
+
+        @enginefacade.reader.async
+        def go1(context):
+            context.session.execute("test1")
+            go2(context)
+
+        @enginefacade.reader.allow_async
+        def go2(context):
+            context.session.execute("test2")
+
+        go1(context)
+
+        with self._assert_engines() as engines:
+            with self._assert_makers(engines) as makers:
+                with self._assert_async_reader_session(makers) as session:
+                    session.execute("test1")
+                    session.execute("test2")
+
+    def test_reader_allow_async_nested_in_reader(self):
+        context = oslo_context.RequestContext()
+
+        @enginefacade.reader.reader
+        def go1(context):
+            context.session.execute("test1")
+            go2(context)
+
+        @enginefacade.reader.allow_async
+        def go2(context):
+            context.session.execute("test2")
+
+        go1(context)
+
+        with self._assert_engines() as engines:
+            with self._assert_makers(engines) as makers:
+                with self._assert_reader_session(makers) as session:
+                    session.execute("test1")
+                    session.execute("test2")
+
+    def test_reader_allow_async_is_reader_by_default(self):
+        context = oslo_context.RequestContext()
+
+        @enginefacade.reader.allow_async
+        def go1(context):
+            context.session.execute("test1")
+
+        go1(context)
+
+        with self._assert_engines() as engines:
+            with self._assert_makers(engines) as makers:
+                with self._assert_reader_session(makers) as session:
+                    session.execute("test1")
+
     def test_writer_nested_in_async_reader_raises(self):
         context = oslo_context.RequestContext()