Add support for pickling enginefacade context objects

In order to pickle a context object which refers
to an active _TransactionContextTLocal, we have to
use __reduce__ so that no database, engine, or session
state is pickled, and only a constructor remains.  Upon
unpickling, the constructor will install a new, empty
_TransactionContextTLocal on the new context, as is
appropriate since this is a totally different context object
than the original one.

Change-Id: Ia50cf10ff91a013ee24773c095a3df69ae06847b

2 files changed, 20 insertions, 1 deletions

diff --git a/oslo_db/sqlalchemy/enginefacade.py b/oslo_db/sqlalchemy/enginefacade.py
index cf7f149..8b241cf 100644
--- a/oslo_db/sqlalchemy/enginefacade.py
+++ b/oslo_db/sqlalchemy/enginefacade.py
@@ -569,6 +569,9 @@ class _TransactionContextTLocal(threading.local):
     def __deepcopy__(self, memo):
         return self
 
+    def __reduce__(self):
+        return _TransactionContextTLocal, ()
+
 
 class _TransactionContextManager(object):
     """Provide context-management and decorator patterns for transactions.
diff --git a/oslo_db/tests/sqlalchemy/test_enginefacade.py b/oslo_db/tests/sqlalchemy/test_enginefacade.py
index 6893a93..4a2517e 100644
--- a/oslo_db/tests/sqlalchemy/test_enginefacade.py
+++ b/oslo_db/tests/sqlalchemy/test_enginefacade.py
@@ -13,6 +13,7 @@
 import collections
 import contextlib
 import copy
+import pickle
 import warnings
 
 import mock
@@ -986,7 +987,7 @@ class LegacyIntegrationtest(test_base.DbTestCase):
 
 
 class ThreadingTest(test_base.DbTestCase):
-    """Test copying on new threads using real connections and sessions."""
+    """Test copy/pickle on new threads using real connections and sessions."""
 
     def _assert_ctx_connection(self, context, connection):
         self.assertIs(context.connection, connection)
@@ -1162,6 +1163,21 @@ class ThreadingTest(test_base.DbTestCase):
         with self._patch_thread_ident():
             go_one(context)
 
+    def test_contexts_picklable(self):
+        context = oslo_context.RequestContext()
+
+        with enginefacade.writer.using(context) as session:
+            self._assert_ctx_session(context, session)
+
+            pickled = pickle.dumps(context)
+
+            unpickled = pickle.loads(pickled)
+
+            with enginefacade.writer.using(unpickled) as session2:
+                self._assert_ctx_session(unpickled, session2)
+
+                assert session is not session2
+
 
 class LiveFacadeTest(test_base.DbTestCase):
     """test using live SQL with test-provisioned databases.