Check context

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,

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):

tests/unit/test_client.py

@@ -95,10 +95,19 @@ def test__http():
             client._http
 
     @staticmethod
-    def test__context():
+    def test_context():
         with patch_credentials("testing"):
             client = client_module.Client()
 
         with client.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__()