NDB: Implement Client.context() (#7013)

As far as the public API is concerned, ``Client.context()`` replaces ``ndb_context``, which has been renamed to ``state_context`` and removed from the public API. ``_api.stub()`` has been refactored to store the Datastore stub in the current context.

Author: Chris Rossi

ndb/docs/client.rst

@@ -0,0 +1,7 @@
+######
+Client
+######
+
+.. automodule:: google.cloud.ndb.client
+ :members:
+ :show-inheritance:

ndb/docs/conf.py

@@ -207,6 +207,7 @@
 # Example configuration for intersphinx: refer to the Python standard library.
 intersphinx_mapping = {
  "python": ("https://docs.python.org/", None),
+ "google-auth": ("https://google-auth.readthedocs.io/en/latest/", None),
  "google-cloud-datastore": (
  "https://googleapis.github.io/google-cloud-python/latest/",
  None,

ndb/docs/context.rst

@@ -6,6 +6,7 @@
  :hidden:
  :maxdepth: 2
 
+ client
  key
  model
  query
@@ -16,7 +17,6 @@
  blobstore
  metadata
  stats
- context
 
 .. automodule:: google.cloud.ndb
  :no-members:

ndb/docs/index.rst

@@ -67,7 +67,6 @@
  "ModelAdapter",
  "ModelAttribute",
  "ModelKey",
- "ndb_context",
  "non_transactional",
  "PickleProperty",
  "Property",
@@ -205,7 +204,6 @@
 from google.cloud.ndb.query import QueryIterator
 from google.cloud.ndb.query import QueryOptions
 from google.cloud.ndb.query import RepeatedStructuredPropertyPredicate
-from google.cloud.ndb._runstate import ndb_context
 from google.cloud.ndb.tasklets import add_flow_exception
 from google.cloud.ndb.tasklets import Future
 from google.cloud.ndb.tasklets import get_context

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

@@ -20,22 +20,30 @@
 from google.cloud import _http
 from google.cloud.datastore_v1.proto import datastore_pb2_grpc
 
+from google.cloud.ndb import _runstate
 
-def stub(client):
- """Get a stub for the `Google Datastore` API.
 
- Arguments:
- client (:class:`~client.Client`): An NDB client instance.
+def stub():
+ """Get the stub for the `Google Datastore` API.
+
+ Gets the stub from the current context, creating one if there isn't one
+ already.
 
  Returns:
  :class:`~google.cloud.datastore_v1.proto.datastore_pb2_grpc.DatastoreStub`:
  The stub instance.
  """
- if client.secure:
- channel = _helpers.make_secure_channel(
- client._credentials, _http.DEFAULT_USER_AGENT, client.host
- )
- else:
- channel = grpc.insecure_channel(client.host)
- stub = datastore_pb2_grpc.DatastoreStub(channel)
- return stub
+ state = _runstate.current()
+
+ if state.stub is None:
+ client = state.client
+ if client.secure:
+ channel = _helpers.make_secure_channel(
+ client._credentials, _http.DEFAULT_USER_AGENT, client.host
+ )
+ else:
+ channel = grpc.insecure_channel(client.host)
+
+ state.stub = datastore_pb2_grpc.DatastoreStub(channel)
+
+ return state.stub

ndb/src/google/cloud/ndb/_api.py

@@ -21,7 +21,10 @@
 
 
 class State:
- eventloop = None
+ def __init__(self, client):
+ self.client = client
+ self.eventloop = None
+ self.stub = None
 
 
 class LocalStates(threading.local):
@@ -47,34 +50,13 @@ def current(self):
 
 
 @contextlib.contextmanager
-def ndb_context():
+def state_context(client):
  """Establish a context for a set of NDB calls.
 
- This function provides a context manager which establishes the runtime
- state for using NDB.
-
- For example:
-
- .. code-block:: python
-
- from google.cloud.ndb import ndb_context
-
- with ndb_context():
- # Use NDB for some stuff
- pass
-
- Use of a context is required--NDB can only be used inside a running
- context. The context is used to coordinate an event loop for asynchronous
- API calls, runtime caching policy, and other essential runtime state.
-
- Code within an asynchronous context should be single threaded. Internally,
- a :class:`threading.local` instance is used to track the current event
- loop.
-
- In a web application, it is recommended that a single context be used per
- HTTP request. This can typically be accomplished in a middleware layer.
+ Called from :meth:`google.cloud.ndb.client.Client.context` which has more
+ information.
  """
- state = State()
+ state = State(client)
  states.push(state)
  yield
 
@@ -91,14 +73,14 @@ def current():
  """Get the current context state.
 
  This function should be called within a context established by
- :func:`~google.cloud.ndb.ndb_context`.
+ :meth:`google.cloud.ndb.client.Client.context`.
 
  Returns:
  State: The state for the current context.
 
  Raises:
  .ContextError: If called outside of a context
- established by :func:`~google.cloud.ndb.ndb_context`.
+ established by :meth:`google.cloud.ndb.client.Client.context`.
  """
  state = states.current()
  if state:

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

@@ -21,9 +21,11 @@
 from google.cloud import client as google_client
 from google.cloud.datastore_v1.gapic import datastore_client
 
-DATASTORE_API_HOST = datastore_client.DatastoreClient.SERVICE_ADDRESS.rstrip(
- ":443"
-)
+from google.cloud.ndb import _runstate
+
+DATASTORE_API_HOST = datastore_client.DatastoreClient.SERVICE_ADDRESS.rsplit(
+ ":", 1
+)[0]
 
 
 def _get_gcd_project():
@@ -60,6 +62,9 @@ def _determine_default_project(project=None):
 class Client(google_client.ClientWithProject):
  """An NDB client.
 
+ The NDB client must be created in order to use NDB, and any use of NDB must
+ be within the context of a call to :meth:`context`.
+
  Arguments:
  project (Optional[str]): The project to pass to proxied API methods. If
  not passed, falls back to the default inferred from the
@@ -73,15 +78,45 @@ class Client(google_client.ClientWithProject):
  SCOPE = ("https://www.googleapis.com/auth/datastore",)
  """The scopes required for authenticating as a Cloud Datastore consumer."""
 
- secure = True
- """Whether to use a secure connection for API calls."""
-
  def __init__(self, project=None, namespace=None, credentials=None):
  super(Client, self).__init__(project=project, credentials=credentials)
  self.namespace = namespace
  self.host = os.environ.get(
  environment_vars.GCD_HOST, DATASTORE_API_HOST
  )
+ self.secure = True
+
+ def context(self):
+ """Establish a context for a set of NDB calls.
+
+ This method provides a context manager which establishes the runtime
+ state for using NDB.
+
+ For example:
+
+ .. code-block:: python
+
+ from google.cloud import ndb
+
+ client = ndb.Client()
+ with client.context():
+ # Use NDB for some stuff
+ pass
+
+ Use of a context is required--NDB can only be used inside a running
+ context. The context is used to manage the connection to Google Cloud
+ Datastore, an event loop for asynchronous API calls, runtime caching
+ policy, and other essential runtime state.
+
+ Code within an asynchronous context should be single threaded.
+ Internally, a :class:`threading.local` instance is used to track the
+ current event loop.
+
+ In a web application, it is recommended that a single context be used
+ per HTTP request. This can typically be accomplished in a middleware
+ layer.
+ """
+ return _runstate.state_context(self)
 
  @property
  def _http(self):

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

@@ -38,13 +38,13 @@ class ContextError(Error):
  """Indicates an NDB call being made without a context.
 
  Raised whenever an NDB call is made outside of a context
- established by :func:`~google.cloud.ndb.ndb_context`.
+ established by :meth:`google.cloud.ndb.client.Client.context`.
  """
 
  def __init__(self):
  super(ContextError, self).__init__(
  "No currently running event loop. Asynchronous calls must be made "
- "in context established by google.cloud.ndb.ndb_context."
+ "in context established by google.cloud.ndb.Client.context."
  )
 
 

ndb/src/google/cloud/ndb/exceptions.py

@@ -16,6 +16,7 @@
 
 from google.cloud import _http
 from google.cloud.ndb import _api
+from google.cloud.ndb import _runstate
 
 
 class TestStub:
@@ -30,7 +31,9 @@ def test_secure_channel(datastore_pb2_grpc, _helpers):
  host="thehost",
  spec=("_credentials", "secure", "host"),
  )
- stub = _api.stub(client)
+ with _runstate.state_context(client):
+ stub = _api.stub()
+ assert _api.stub() is stub # one stub per context
  assert stub is datastore_pb2_grpc.DatastoreStub.return_value
  datastore_pb2_grpc.DatastoreStub.assert_called_once_with(channel)
  _helpers.make_secure_channel.assert_called_once_with(
@@ -45,7 +48,8 @@ def test_insecure_channel(datastore_pb2_grpc, grpc):
  client = mock.Mock(
  secure=False, host="thehost", spec=("secure", "host")
  )
- stub = _api.stub(client)
+ with _runstate.state_context(client):
+ stub = _api.stub()
  assert stub is datastore_pb2_grpc.DatastoreStub.return_value
  datastore_pb2_grpc.DatastoreStub.assert_called_once_with(channel)
  grpc.insecure_channel.assert_called_once_with("thehost")