fix: move stub (grpc communication channel) to client

google/cloud/ndb/_datastore_api.py

@@ -17,12 +17,8 @@
 import itertools
 import logging
 
-import grpc
-
-from google.cloud import _helpers
 from google.cloud.datastore import helpers
 from google.cloud.datastore_v1.proto import datastore_pb2
-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
@@ -54,28 +50,7 @@ def stub():
             The stub instance.
     """
     context = context_module.get_context()
-    return context.stub
-
-
-def make_stub(client):
-    """Create the stub for the `Google Datastore` API.
-
-    Args:
-        client (client.Client): The NDB client.
-
-    Returns:
-        :class:`~google.cloud.datastore_v1.proto.datastore_pb2_grpc.DatastoreStub`:
-            The stub instance.
-    """
-    if client.secure:
-        user_agent = client.client_info.to_user_agent()
-        channel = _helpers.make_secure_channel(
-            client._credentials, user_agent, client.host
-        )
-    else:
-        channel = grpc.insecure_channel(client.host)
-
-    return datastore_pb2_grpc.DatastoreStub(channel)
+    return context.client.stub
 
 
 def make_call(rpc_name, request, retries=None, timeout=None):

google/cloud/ndb/client.py

@@ -15,6 +15,7 @@
 """A client for NDB which manages credentials, project, namespace."""
 
 import contextlib
+import grpc
 import os
 import requests
 
@@ -23,10 +24,12 @@
 from google.cloud import _helpers
 from google.cloud import client as google_client
 from google.cloud.datastore_v1.gapic import datastore_client
+from google.cloud.datastore_v1.proto import datastore_pb2_grpc
 
 from google.cloud.ndb import __version__
 from google.cloud.ndb import context as context_module
 
+
 _CLIENT_INFO = client_info.ClientInfo(
     user_agent="google-cloud-ndb/{}".format(__version__)
 )
@@ -114,6 +117,17 @@ def __init__(self, project=None, namespace=None, credentials=None):
                 project=project, credentials=credentials
             )
 
+        if emulator:
+            channel = grpc.insecure_channel(self.host)
+
+        else:
+            user_agent = _CLIENT_INFO.to_user_agent()
+            channel = _helpers.make_secure_channel(
+                self._credentials, user_agent, self.host
+            )
+
+        self.stub = datastore_pb2_grpc.DatastoreStub(channel)
+
     @contextlib.contextmanager
     def context(
         self,

google/cloud/ndb/context.py

@@ -146,7 +146,6 @@ def policy(key):
     [
         "client",
         "eventloop",
-        "stub",
         "batches",
         "commit_batches",
         "transaction",
@@ -179,7 +178,6 @@ def __new__(
         cls,
         client,
         eventloop=None,
-        stub=None,
         batches=None,
         commit_batches=None,
         transaction=None,
@@ -194,14 +192,10 @@ def __new__(
     ):
         # Prevent circular import in Python 2.7
         from google.cloud.ndb import _cache
-        from google.cloud.ndb import _datastore_api
 
         if eventloop is None:
             eventloop = _eventloop.EventLoop()
 
-        if stub is None:
-            stub = _datastore_api.make_stub(client)
-
         if batches is None:
             batches = {}
 
@@ -218,7 +212,6 @@ def __new__(
             cls,
             client=client,
             eventloop=eventloop,
-            stub=stub,
             batches=batches,
             commit_batches=commit_batches,
             transaction=transaction,

google/cloud/ndb/key.py

@@ -147,8 +147,13 @@ class Key(object):
 
         from unittest import mock
         from google.cloud.ndb import context as context_module
-        client = mock.Mock(project="testing", spec=("project",), namespace="")
-        context = context_module.Context(client, stub=mock.Mock(spec=())).use()
+        client = mock.Mock(
+            project="testing",
+            namespace="",
+            stub=mock.Mock(spec=()),
+            spec=("project", "namespace", "stub"),
+        )
+        context = context_module.Context(client).use()
         context.__enter__()
         kind1, id1 = "Parent", "C"
         kind2, id2 = "Child", 42

google/cloud/ndb/model.py

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

tests/conftest.py

@@ -87,11 +87,13 @@ def initialize_environment(request, environ):
 @pytest.fixture
 def context():
     client = mock.Mock(
-        project="testing", namespace=None, spec=("project", "namespace")
+        project="testing",
+        namespace=None,
+        spec=("project", "namespace"),
+        stub=mock.Mock(spec=()),
     )
     context = context_module.Context(
         client,
-        stub=mock.Mock(spec=()),
         eventloop=TestingEventLoop(),
         datastore_policy=True,
         legacy_data=False,

tests/unit/test__datastore_api.py

@@ -46,43 +46,20 @@ def future_result(result):
 
 class TestStub:
     @staticmethod
-    @mock.patch("google.cloud.ndb._datastore_api._helpers")
-    @mock.patch("google.cloud.ndb._datastore_api.datastore_pb2_grpc")
-    def test_secure_channel(datastore_pb2_grpc, _helpers):
-        channel = _helpers.make_secure_channel.return_value
+    def test_it():
         client = mock.Mock(
             _credentials="creds",
             secure=True,
             host="thehost",
-            spec=("_credentials", "secure", "host"),
+            stub=object(),
+            spec=("_credentials", "secure", "host", "stub"),
             client_info=client_info.ClientInfo(
                 user_agent="google-cloud-ndb/{}".format(__version__)
             ),
         )
         context = context_module.Context(client)
         with context.use():
-            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(
-            "creds", client.client_info.to_user_agent(), "thehost"
-        )
-
-    @staticmethod
-    @mock.patch("google.cloud.ndb._datastore_api.grpc")
-    @mock.patch("google.cloud.ndb._datastore_api.datastore_pb2_grpc")
-    def test_insecure_channel(datastore_pb2_grpc, grpc):
-        channel = grpc.insecure_channel.return_value
-        client = mock.Mock(
-            secure=False, host="thehost", spec=("secure", "host")
-        )
-        context = context_module.Context(client)
-        with context.use():
-            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")
+            assert _api.stub() is client.stub
 
 
 class Test_make_call:
@@ -510,10 +487,13 @@ def key_pb(key):
 
 @mock.patch("google.cloud.ndb._datastore_api.datastore_pb2")
 def test__datastore_lookup(datastore_pb2, context):
-    client = mock.Mock(project="theproject", spec=("project",))
-    stub = mock.Mock(spec=("Lookup",))
-    with context.new(client=client, stub=stub).use() as context:
-        context.stub.Lookup = Lookup = mock.Mock(spec=("future",))
+    client = mock.Mock(
+        project="theproject",
+        stub=mock.Mock(spec=("Lookup",)),
+        spec=("project", "stub"),
+    )
+    with context.new(client=client).use() as context:
+        client.stub.Lookup = Lookup = mock.Mock(spec=("future",))
         future = tasklets.Future()
         future.set_result("response")
         Lookup.future.return_value = future
@@ -524,7 +504,7 @@ def test__datastore_lookup(datastore_pb2, context):
         datastore_pb2.LookupRequest.assert_called_once_with(
             project_id="theproject", keys=["foo", "bar"], read_options=None
         )
-        context.stub.Lookup.future.assert_called_once_with(
+        client.stub.Lookup.future.assert_called_once_with(
             datastore_pb2.LookupRequest.return_value,
             timeout=_api._DEFAULT_TIMEOUT,
         )

tests/unit/test_context.py

@@ -37,31 +37,28 @@ def test___all__():
 class TestContext:
     def _make_one(self, **kwargs):
         client = mock.Mock(
-            namespace=None, project="testing", spec=("namespace", "project")
+            namespace=None,
+            project="testing",
+            spec=("namespace", "project"),
+            stub=mock.Mock(spec=()),
         )
-        stub = mock.Mock(spec=())
-        return context_module.Context(client, stub=stub, **kwargs)
+        return context_module.Context(client, **kwargs)
 
-    @mock.patch("google.cloud.ndb._datastore_api.make_stub")
-    def test_constructor_defaults(self, make_stub):
+    def test_constructor_defaults(self):
         context = context_module.Context("client")
         assert context.client == "client"
-        assert context.stub is make_stub.return_value
-        make_stub.assert_called_once_with("client")
         assert isinstance(context.eventloop, _eventloop.EventLoop)
         assert context.batches == {}
         assert context.transaction is None
 
     def test_constructor_overrides(self):
         context = context_module.Context(
             client="client",
-            stub="stub",
             eventloop="eventloop",
             batches="batches",
             transaction="transaction",
         )
         assert context.client == "client"
-        assert context.stub == "stub"
         assert context.eventloop == "eventloop"
         assert context.batches == "batches"
         assert context.transaction == "transaction"