Merge pull request #30 from chrisrossi/context-changes-in-tasklets

Handle contexts being changed internally to a tasklet.

Author: Chris Rossi

packages/google-cloud-ndb/src/google/cloud/ndb/__init__.py

@@ -125,6 +125,7 @@
 from google.cloud.ndb.context import AutoBatcher
 from google.cloud.ndb.context import Context
 from google.cloud.ndb.context import ContextOptions
+from google.cloud.ndb.context import get_context
 from google.cloud.ndb.context import TransactionOptions
 from google.cloud.ndb._datastore_api import EVENTUAL
 from google.cloud.ndb._datastore_api import EVENTUAL_CONSISTENCY
@@ -206,7 +207,6 @@
 from google.cloud.ndb.query import RepeatedStructuredPropertyPredicate
 from google.cloud.ndb.tasklets import add_flow_exception
 from google.cloud.ndb.tasklets import Future
-from google.cloud.ndb.tasklets import get_context
 from google.cloud.ndb.tasklets import make_context
 from google.cloud.ndb.tasklets import make_default_context
 from google.cloud.ndb.tasklets import QueueFuture

packages/google-cloud-ndb/src/google/cloud/ndb/_datastore_api.py

@@ -24,8 +24,8 @@
 from google.cloud.datastore_v1.proto import datastore_pb2_grpc
 from google.cloud.datastore_v1.proto import entity_pb2
 
+from google.cloud.ndb import context as context_module
 from google.cloud.ndb import _eventloop
-from google.cloud.ndb import _runstate
 from google.cloud.ndb import tasklets
 
 EVENTUAL = datastore_pb2.ReadOptions.EVENTUAL
@@ -42,8 +42,8 @@ def stub():
         :class:`~google.cloud.datastore_v1.proto.datastore_pb2_grpc.DatastoreStub`:
             The stub instance.
     """
-    state = _runstate.current()
-    return state.stub
+    context = context_module.get_context()
+    return context.stub
 
 
 def make_stub(client):
@@ -126,7 +126,7 @@ def lookup(key, **options):
 def _get_batch(batch_cls, options):
     """Gets a data structure for storing batched calls to Datastore Lookup.
 
-    The batch data structure is stored in the current run state. If there is
+    The batch data structure is stored in the current context. If there is
     not already a batch started, a new structure is created and an idle
     callback is added to the current event loop which will eventually perform
     the batch look up.
@@ -141,10 +141,10 @@ def _get_batch(batch_cls, options):
     Returns:
         batch_cls: An instance of the batch class.
     """
-    state = _runstate.current()
-    batches = state.batches.get(batch_cls)
+    context = context_module.get_context()
+    batches = context.batches.get(batch_cls)
     if batches is None:
-        state.batches[batch_cls] = batches = {}
+        context.batches[batch_cls] = batches = {}
 
     options_key = tuple(sorted(options.items()))
     batch = batches.get(options_key)
@@ -266,7 +266,7 @@ def _datastore_lookup(keys, read_options):
     Returns:
         RemoteCall: Future object for eventual result of lookup.
     """
-    client = _runstate.current().client
+    client = context_module.get_context().client
     request = datastore_pb2.LookupRequest(
         project_id=client.project,
         keys=[key for key in keys],
@@ -323,8 +323,8 @@ def _get_transaction(options):
     Returns:
         Union[bytes, NoneType]: The transaction identifier, or :data:`None`.
     """
-    state = _runstate.current()
-    return options.get("transaction", state.transaction)
+    context = context_module.get_context()
+    return options.get("transaction", context.transaction)
 
 
 def put(entity_pb, **options):
@@ -445,7 +445,7 @@ def _datastore_commit(mutations, transaction):
     else:
         mode = datastore_pb2.CommitRequest.TRANSACTIONAL
 
-    client = _runstate.current().client
+    client = context_module.get_context().client
     request = datastore_pb2.CommitRequest(
         project_id=client.project,
         mode=mode,

packages/google-cloud-ndb/src/google/cloud/ndb/_eventloop.py

@@ -21,7 +21,7 @@
 import uuid
 import time
 
-from google.cloud.ndb import _runstate
+from google.cloud.ndb import context as context_module
 
 __all__ = [
     "add_idle",
@@ -369,7 +369,7 @@ def get_event_loop():
     Returns:
         EventLoop: The event loop for the current context.
     """
-    context = _runstate.current()
+    context = context_module.get_context()
     return context.eventloop
 
 

packages/google-cloud-ndb/src/google/cloud/ndb/_runstate.py

@@ -119,7 +119,7 @@ def context(self):
         layer.
         """
         context = context_module.Context(self)
-        with context:
+        with context.use():
             yield context
 
         # Finish up any work left to do on the event loop

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

@@ -15,21 +15,60 @@
 """Context for currently running tasks and transactions."""
 
 import collections
+import contextlib
+import threading
 
 from google.cloud.ndb import _datastore_api
 from google.cloud.ndb import _eventloop
 from google.cloud.ndb import exceptions
-from google.cloud.ndb import _runstate
 
 
-__all__ = ["AutoBatcher", "Context", "ContextOptions", "TransactionOptions"]
+__all__ = [
+    "AutoBatcher",
+    "Context",
+    "ContextOptions",
+    "get_context",
+    "TransactionOptions",
+]
 
 
 _ContextTuple = collections.namedtuple(
     "_ContextTuple", ["client", "eventloop", "stub", "batches", "transaction"]
 )
 
 
+class _LocalState(threading.local):
+    """Thread local state."""
+
+    __slots__ = ("context",)
+
+    def __init__(self):
+        self.context = None
+
+
+_state = _LocalState()
+
+
+def get_context():
+    """Get the current context.
+
+    This function should be called within a context established by
+    :meth:`google.cloud.ndb.client.Client.context`.
+
+    Returns:
+        Context: The current context.
+
+    Raises:
+        .ContextError: If called outside of a context
+            established by :meth:`google.cloud.ndb.client.Client.context`.
+    """
+    context = _state.context
+    if context:
+        return context
+
+    raise exceptions.ContextError()
+
+
 class _Context(_ContextTuple):
     """Current runtime state.
 
@@ -39,10 +78,6 @@ class _Context(_ContextTuple):
     loop. A new context can be derived from an existing context using
     :meth:`new`.
 
-    ``_Context`` instances can be used as context managers which push
-    themselves onto the thread local stack in ``_runstate`` and then pop
-    themselves back off on exit.
-
     :class:`Context` is a subclass of :class:`_Context` which provides
     only publicly facing interface. The use of two classes is only to provide a
     distinction between public and private API.
@@ -82,17 +117,20 @@ def new(self, **kwargs):
         state.update(kwargs)
         return type(self)(**state)
 
-    def __enter__(self):
-        _runstate.contexts.push(self)
-        return self
+    @contextlib.contextmanager
+    def use(self):
+        """Use this context as the current context.
 
-    def __exit__(self, *exc_info):
-        popped = _runstate.contexts.pop()
-
-        # If we've done this right, this will never happen. Including this
-        # check in an abundance of caution.
-        if popped is not self:
-            raise RuntimeError("Contexts stack is corrupted")
+        This method returns a context manager for use with the ``with``
+        statement. Code inside the ``with`` context will see this context as
+        the current context.
+        """
+        prev_context = _state.context
+        _state.context = self
+        try:
+            yield self
+        finally:
+            _state.context = prev_context
 
 
 class Context(_Context):

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

@@ -91,9 +91,9 @@
 from google.cloud.datastore import key as _key_module
 import google.cloud.datastore
 
+from google.cloud.ndb import context as context_module
 from google.cloud.ndb import _datastore_api
 from google.cloud.ndb import exceptions
-from google.cloud.ndb import _runstate
 from google.cloud.ndb import tasklets
 
 
@@ -137,7 +137,7 @@ class Key:
         from unittest import mock
         from google.cloud.ndb import context as context_module
         client = mock.Mock(project="testing", spec=("project",))
-        context = context_module.Context(client, stub=mock.Mock(spec=()))
+        context = context_module.Context(client, stub=mock.Mock(spec=())).use()
         context.__enter__()
         kind1, id1 = "Parent", "C"
         kind2, id2 = "Child", 42
@@ -816,7 +816,7 @@ def _project_from_app(app, allow_empty=False):
     if app is None:
         if allow_empty:
             return None
-        client = _runstate.current().client
+        client = context_module.get_context().client
         app = client.project
 
     # NOTE: This is the same behavior as in the helper

packages/google-cloud-ndb/src/google/cloud/ndb/key.py

@@ -21,7 +21,7 @@
     from google.cloud.ndb import context as context_module
 
     client = mock.Mock(project="testing", spec=("project",))
-    context = context_module.Context(client, stub=mock.Mock(spec=()))
+    context = context_module.Context(client, stub=mock.Mock(spec=())).use()
     context.__enter__()
 
 .. testcleanup:: *

packages/google-cloud-ndb/src/google/cloud/ndb/model.py

@@ -22,13 +22,12 @@
 
 import grpc
 
+from google.cloud.ndb import context as context_module
 from google.cloud.ndb import _eventloop
-from google.cloud.ndb import _runstate
 
 __all__ = [
     "add_flow_exception",
     "Future",
-    "get_context",
     "make_context",
     "make_default_context",
     "QueueFuture",
@@ -248,14 +247,17 @@ def __init__(self, generator, context, info="Unknown"):
     def _advance_tasklet(self, send_value=None, error=None):
         """Advance a tasklet one step by sending in a value or error."""
         try:
-            with self.context:
+            with self.context.use():
                 # Send the next value or exception into the generator
                 if error:
                     self.generator.throw(type(error), error)
 
                 # send_value will be None if this is the first time
                 yielded = self.generator.send(send_value)
 
+                # Context may have changed in tasklet
+                self.context = context_module.get_context()
+
         except StopIteration as stop:
             # Generator has signalled exit, get the return value. This tasklet
             # has finished.
@@ -380,7 +382,7 @@ def tasklet_wrapper(*args, **kwargs):
         # and create a future object and set the result to the function's
         # return value so that from the user perspective there is no problem.
         # This permissive behavior is inherited from legacy NDB.
-        context = _runstate.current()
+        context = context_module.get_context()
 
         try:
             returned = wrapped(*args, **kwargs)
@@ -470,10 +472,6 @@ def add_flow_exception(*args, **kwargs):
     raise NotImplementedError
 
 
-def get_context(*args, **kwargs):
-    raise NotImplementedError
-
-
 def make_context(*args, **kwargs):
     raise NotImplementedError
 

packages/google-cloud-ndb/src/google/cloud/ndb/tasklets.py

@@ -25,7 +25,6 @@
 from google.cloud import environment_vars
 from google.cloud.ndb import context as context_module
 from google.cloud.ndb import model
-from google.cloud.ndb import _runstate
 
 import pytest
 
@@ -45,7 +44,6 @@ def reset_state(environ):
     yield
     model.Property._FIND_METHODS_CACHE.clear()
     model.Model._kind_map.clear()
-    del _runstate.contexts.stack[:]
 
 
 @pytest.fixture
@@ -84,5 +82,7 @@ def context():
 
 @pytest.fixture
 def in_context(context):
-    with context:
+    assert not context_module._state.context
+    with context.use():
         yield context
+    assert not context_module._state.context