From d91be8e8e6c45bb19c4b46917d0a73a836565035 Mon Sep 17 00:00:00 2001 From: Chris Rossi Date: Wed, 11 Mar 2020 10:46:33 -0400 Subject: [PATCH] fix: move stub (grpc communication channel) to client This brings our practice in line with `google.cloud.datastore`, which also creates one channel per client. This works around a resource leak issue by not requiring the channel to clean up after itself properly in normal usage. The root cause of that issue seems to lie somewhere in `google.auth`, which is where I will follow up. Fixes #343 --- google/cloud/ndb/_datastore_api.py | 27 +----------------- google/cloud/ndb/client.py | 14 ++++++++++ google/cloud/ndb/context.py | 7 ----- google/cloud/ndb/key.py | 9 ++++-- google/cloud/ndb/model.py | 9 ++++-- tests/conftest.py | 6 ++-- tests/unit/test__datastore_api.py | 44 ++++++++---------------------- tests/unit/test_context.py | 15 ++++------ 8 files changed, 51 insertions(+), 80 deletions(-) diff --git a/google/cloud/ndb/_datastore_api.py b/google/cloud/ndb/_datastore_api.py index ec995c6b..92a6426b 100644 --- a/google/cloud/ndb/_datastore_api.py +++ b/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): diff --git a/google/cloud/ndb/client.py b/google/cloud/ndb/client.py index 56067fb2..fa9ad890 100644 --- a/google/cloud/ndb/client.py +++ b/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, diff --git a/google/cloud/ndb/context.py b/google/cloud/ndb/context.py index 94cfe640..a082bd4e 100644 --- a/google/cloud/ndb/context.py +++ b/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, diff --git a/google/cloud/ndb/key.py b/google/cloud/ndb/key.py index 228af570..fd819256 100644 --- a/google/cloud/ndb/key.py +++ b/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 diff --git a/google/cloud/ndb/model.py b/google/cloud/ndb/model.py index c2cd05bd..3605ade4 100644 --- a/google/cloud/ndb/model.py +++ b/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:: * diff --git a/tests/conftest.py b/tests/conftest.py index 05dc29f0..c6a7db1c 100644 --- a/tests/conftest.py +++ b/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, diff --git a/tests/unit/test__datastore_api.py b/tests/unit/test__datastore_api.py index 7183d528..7cab0033 100644 --- a/tests/unit/test__datastore_api.py +++ b/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, ) diff --git a/tests/unit/test_context.py b/tests/unit/test_context.py index a69672de..32f429d0 100644 --- a/tests/unit/test_context.py +++ b/tests/unit/test_context.py @@ -37,17 +37,16 @@ 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 @@ -55,13 +54,11 @@ def test_constructor_defaults(self, make_stub): 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"