Check context (#211)

When creating a top level context, check to make sure there isn't
already a running context for this thread. This is an attempt to root
a possible weird concurrency problem. Related to #182.

Author: Chris Rossi

packages/google-cloud-ndb/google/cloud/ndb/client.py

@@ -161,6 +161,10 @@ def context(
             legacy_data (bool): Set to ``True`` (the default) to write data in
                 a way that can be read by the legacy version of NDB.
         """
+        context = context_module.get_context(False)
+        if context is not None:
+            raise RuntimeError("Context is already created for this thread.")
+
         context = context_module.Context(
             self,
             cache_policy=cache_policy,

packages/google-cloud-ndb/google/cloud/ndb/context.py

@@ -45,24 +45,32 @@ def __init__(self):
 _state = _LocalState()
 
 
-def get_context():
+def get_context(raise_context_error=True):
     """Get the current context.
 
     This function should be called within a context established by
     :meth:`google.cloud.ndb.client.Client.context`.
 
+    Args:
+        raise_context_error (bool): If set to :data:`True`, will raise an
+            exception if called outside of a context. Set this to :data:`False`
+            in order to have it just return :data:`None` if called outside of a
+            context. Default: :data:`True`
+
     Returns:
         Context: The current context.
 
     Raises:
         .ContextError: If called outside of a context
-            established by :meth:`google.cloud.ndb.client.Client.context`.
+            established by :meth:`google.cloud.ndb.client.Client.context` and
+            ``raise_context_error`` is :data:`True`.
     """
     context = _state.context
     if context:
         return context
 
-    raise exceptions.ContextError()
+    if raise_context_error:
+        raise exceptions.ContextError()
 
 
 def _default_policy(attr_name, value_type):

packages/google-cloud-ndb/tests/unit/test_client.py

@@ -104,6 +104,15 @@ def test_context():
             context = context_module.get_context()
             assert context.client is client
 
+    @staticmethod
+    def test_context_double_jeopardy():
+        with patch_credentials("testing"):
+            client = client_module.Client()
+
+        with client.context():
+            with pytest.raises(RuntimeError):
+                client.context().__enter__()
+
     @staticmethod
     def test_context_unfinished_business():
         """Regression test for #213.