From 7e31ba4eff9885d868cdfc52bda913373b47f410 Mon Sep 17 00:00:00 2001 From: Danny Hermes Date: Thu, 23 Feb 2017 14:25:45 -0800 Subject: [PATCH] Add basic helpers needed for GAPIC client in datastore. Includes - ``use_gax`` argument in Client constructor - ``make_datastore_api`` helper to make a gRPC channel with the correct credentials - A basic ``HTTPDatastoreAPI`` class to act as the equivalent to the GAPIC ``DatastoreClient`` - A lazy property on ``Client`` that will hold the API object. --- datastore/google/cloud/datastore/_gax.py | 19 ++++++ datastore/google/cloud/datastore/_http.py | 13 ++++ datastore/google/cloud/datastore/client.py | 36 ++++++++++- datastore/unit_tests/test__gax.py | 31 +++++++++ datastore/unit_tests/test__http.py | 17 +++++ datastore/unit_tests/test_client.py | 74 ++++++++++++++++++++-- 6 files changed, 183 insertions(+), 7 deletions(-) diff --git a/datastore/google/cloud/datastore/_gax.py b/datastore/google/cloud/datastore/_gax.py index 58500301c950..1bf7b21fbe5f 100644 --- a/datastore/google/cloud/datastore/_gax.py +++ b/datastore/google/cloud/datastore/_gax.py @@ -17,12 +17,15 @@ import contextlib +from google.cloud.gapic.datastore.v1 import datastore_client from google.cloud.proto.datastore.v1 import datastore_pb2_grpc from google.gax.utils import metrics from grpc import StatusCode from google.cloud._helpers import make_insecure_stub +from google.cloud._helpers import make_secure_channel from google.cloud._helpers import make_secure_stub +from google.cloud._http import DEFAULT_USER_AGENT from google.cloud import exceptions from google.cloud.datastore import __version__ @@ -204,3 +207,19 @@ def allocate_ids(self, project, request_pb): request_pb.project_id = project with _grpc_catch_rendezvous(): return self._stub.AllocateIds(request_pb) + + +def make_datastore_api(client): + """Create an instance of the GAPIC Datastore API. + + :type client: :class:`~google.cloud.datastore.client.Client` + :param client: The client that holds configuration details. + + :rtype: :class:`.datastore.v1.datastore_client.DatastoreClient` + :returns: A datastore API instance with the proper credentials. + """ + channel = make_secure_channel( + client._credentials, DEFAULT_USER_AGENT, + datastore_client.DatastoreClient.SERVICE_ADDRESS) + return datastore_client.DatastoreClient( + channel=channel, lib_name='gccl', lib_version=__version__) diff --git a/datastore/google/cloud/datastore/_http.py b/datastore/google/cloud/datastore/_http.py index a5231431820d..3e82ffbb6e16 100644 --- a/datastore/google/cloud/datastore/_http.py +++ b/datastore/google/cloud/datastore/_http.py @@ -467,6 +467,19 @@ def allocate_ids(self, project, key_pbs): return self._datastore_api.allocate_ids(project, request) +class HTTPDatastoreAPI(object): + """An API object that sends proto-over-HTTP requests. + + Intended to provide the same methods as the GAPIC ``DatastoreClient``. + + :type client: :class:`~google.cloud.datastore.client.Client` + :param client: The client that provides configuration. + """ + + def __init__(self, client): + self.client = client + + def _set_read_options(request, eventual, transaction_id): """Validate rules for read options, and assign to the request. diff --git a/datastore/google/cloud/datastore/client.py b/datastore/google/cloud/datastore/client.py index 438996ceff36..ce924edda183 100644 --- a/datastore/google/cloud/datastore/client.py +++ b/datastore/google/cloud/datastore/client.py @@ -19,19 +19,30 @@ from google.cloud._helpers import ( _determine_default_project as _base_default_project) from google.cloud.client import ClientWithProject +from google.cloud.environment_vars import DISABLE_GRPC +from google.cloud.environment_vars import GCD_DATASET + from google.cloud.datastore._http import Connection +from google.cloud.datastore._http import HTTPDatastoreAPI from google.cloud.datastore import helpers from google.cloud.datastore.batch import Batch from google.cloud.datastore.entity import Entity from google.cloud.datastore.key import Key from google.cloud.datastore.query import Query from google.cloud.datastore.transaction import Transaction -from google.cloud.environment_vars import GCD_DATASET +try: + from google.cloud.datastore._gax import make_datastore_api + _HAVE_GRPC = True +except ImportError: # pragma: NO COVER + make_datastore_api = None + _HAVE_GRPC = False _MAX_LOOPS = 128 """Maximum number of iterations to wait for deferred keys.""" +_USE_GAX = _HAVE_GRPC and not os.getenv(DISABLE_GRPC, False) + def _get_gcd_project(): """Gets the GCD application ID if it can be inferred.""" @@ -169,24 +180,45 @@ class Client(ClientWithProject): :meth:`~httplib2.Http.request`. If not passed, an ``http`` object is created that is bound to the ``credentials`` for the current object. + + :type use_gax: bool + :param use_gax: (Optional) Explicitly specifies whether + to use the gRPC transport (via GAX) or HTTP. If unset, + falls back to the ``GOOGLE_CLOUD_DISABLE_GRPC`` environment + variable. """ SCOPE = ('https://www.googleapis.com/auth/datastore',) """The scopes required for authenticating as a Cloud Datastore consumer.""" def __init__(self, project=None, namespace=None, - credentials=None, http=None): + credentials=None, http=None, use_gax=None): super(Client, self).__init__( project=project, credentials=credentials, http=http) self._connection = Connection(self) self.namespace = namespace self._batch_stack = _LocalStack() + self._datastore_api_internal = None + if use_gax is None: + self._use_gax = _USE_GAX + else: + self._use_gax = use_gax @staticmethod def _determine_default(project): """Helper: override default project detection.""" return _determine_default_project(project) + @property + def _datastore_api(self): + """Getter for a wrapped API object.""" + if self._datastore_api_internal is None: + if self._use_gax: + self._datastore_api_internal = make_datastore_api(self) + else: + self._datastore_api_internal = HTTPDatastoreAPI(self) + return self._datastore_api_internal + def _push_batch(self, batch): """Push a batch/transaction onto our stack. diff --git a/datastore/unit_tests/test__gax.py b/datastore/unit_tests/test__gax.py index 419c85cc9fdc..aea60801cfe5 100644 --- a/datastore/unit_tests/test__gax.py +++ b/datastore/unit_tests/test__gax.py @@ -269,6 +269,37 @@ def test_allocate_ids(self): [(request_pb, 'AllocateIds')]) +@unittest.skipUnless(_HAVE_GRPC, 'No gRPC') +class Test_make_datastore_api(unittest.TestCase): + + def _call_fut(self, client): + from google.cloud.datastore._gax import make_datastore_api + + return make_datastore_api(client) + + @mock.patch( + 'google.cloud.gapic.datastore.v1.datastore_client.DatastoreClient', + SERVICE_ADDRESS='datastore.mock.mock', + return_value=mock.sentinel.ds_client) + @mock.patch('google.cloud.datastore._gax.make_secure_channel', + return_value=mock.sentinel.channel) + def test_it(self, make_chan, mock_klass): + from google.cloud._http import DEFAULT_USER_AGENT + from google.cloud.datastore import __version__ + + client = mock.Mock( + _credentials=mock.sentinel.credentials, spec=['_credentials']) + ds_api = self._call_fut(client) + self.assertIs(ds_api, mock.sentinel.ds_client) + + make_chan.assert_called_once_with( + mock.sentinel.credentials, DEFAULT_USER_AGENT, + mock_klass.SERVICE_ADDRESS) + mock_klass.assert_called_once_with( + channel=mock.sentinel.channel, lib_name='gccl', + lib_version=__version__) + + class _GRPCStub(object): def __init__(self, return_val=None, side_effect=Exception): diff --git a/datastore/unit_tests/test__http.py b/datastore/unit_tests/test__http.py index 23de36cdc2b1..a924fcf40316 100644 --- a/datastore/unit_tests/test__http.py +++ b/datastore/unit_tests/test__http.py @@ -914,6 +914,23 @@ def test_allocate_ids_non_empty(self): self.assertEqual(key_before, key_after) +class TestHTTPDatastoreAPI(unittest.TestCase): + + @staticmethod + def _get_target_class(): + from google.cloud.datastore._http import HTTPDatastoreAPI + + return HTTPDatastoreAPI + + def _make_one(self, *args, **kwargs): + return self._get_target_class()(*args, **kwargs) + + def test_constructor(self): + client = object() + ds_api = self._make_one(client) + self.assertIs(ds_api.client, client) + + class Http(object): _called_with = None diff --git a/datastore/unit_tests/test_client.py b/datastore/unit_tests/test_client.py index 41cf19013a14..63c2e0adebeb 100644 --- a/datastore/unit_tests/test_client.py +++ b/datastore/unit_tests/test_client.py @@ -138,13 +138,14 @@ def _get_target_class(): return Client def _make_one(self, project=PROJECT, namespace=None, - credentials=None, http=None): + credentials=None, http=None, use_gax=None): return self._get_target_class()(project=project, namespace=namespace, credentials=credentials, - http=http) + http=http, + use_gax=use_gax) - def test_ctor_w_project_no_environ(self): + def test_constructor_w_project_no_environ(self): # Some environments (e.g. AppVeyor CI) run in GCE, so # this test would fail artificially. patch = mock.patch( @@ -153,7 +154,7 @@ def test_ctor_w_project_no_environ(self): with patch: self.assertRaises(EnvironmentError, self._make_one, None) - def test_ctor_w_implicit_inputs(self): + def test_constructor_w_implicit_inputs(self): OTHER = 'other' creds = _make_credentials() default_called = [] @@ -183,7 +184,7 @@ def fallback_mock(project): self.assertIsNone(client.current_transaction) self.assertEqual(default_called, [None]) - def test_ctor_w_explicit_inputs(self): + def test_constructor_w_explicit_inputs(self): OTHER = 'other' NAMESPACE = 'namespace' creds = _make_credentials() @@ -200,6 +201,69 @@ def test_ctor_w_explicit_inputs(self): self.assertIsNone(client.current_batch) self.assertEqual(list(client._batch_stack), []) + def test_constructor_use_gax_default(self): + import google.cloud.datastore.client as MUT + + project = 'PROJECT' + creds = _make_credentials() + http = object() + + with mock.patch.object(MUT, '_USE_GAX', new=True): + client1 = self._make_one( + project=project, credentials=creds, http=http) + self.assertTrue(client1._use_gax) + # Explicitly over-ride the environment. + client2 = self._make_one( + project=project, credentials=creds, http=http, + use_gax=False) + self.assertFalse(client2._use_gax) + + with mock.patch.object(MUT, '_USE_GAX', new=False): + client3 = self._make_one( + project=project, credentials=creds, http=http) + self.assertFalse(client3._use_gax) + # Explicitly over-ride the environment. + client4 = self._make_one( + project=project, credentials=creds, http=http, + use_gax=True) + self.assertTrue(client4._use_gax) + + def test__datastore_api_property_gax(self): + client = self._make_one( + project='prahj-ekt', credentials=_make_credentials(), + http=object(), use_gax=True) + + self.assertIsNone(client._datastore_api_internal) + patch = mock.patch( + 'google.cloud.datastore.client.make_datastore_api', + return_value=mock.sentinel.ds_api) + with patch as make_api: + ds_api = client._datastore_api + self.assertIs(ds_api, mock.sentinel.ds_api) + make_api.assert_called_once_with(client) + self.assertIs( + client._datastore_api_internal, mock.sentinel.ds_api) + # Make sure the cached value is used. + self.assertEqual(make_api.call_count, 1) + self.assertIs( + client._datastore_api, mock.sentinel.ds_api) + self.assertEqual(make_api.call_count, 1) + + def test__datastore_api_property_http(self): + from google.cloud.datastore._http import HTTPDatastoreAPI + + client = self._make_one( + project='prahj-ekt', credentials=_make_credentials(), + http=object(), use_gax=False) + + self.assertIsNone(client._datastore_api_internal) + ds_api = client._datastore_api + self.assertIsInstance(ds_api, HTTPDatastoreAPI) + self.assertIs(ds_api.client, client) + # Make sure the cached value is used. + self.assertIs(client._datastore_api_internal, ds_api) + self.assertIs(client._datastore_api, ds_api) + def test__push_batch_and__pop_batch(self): creds = _make_credentials() client = self._make_one(credentials=creds)