Merge pull request googleapis#3105 from dhermes/fix-2746

Add basic helpers needed for GAPIC client in datastore.

Author: dhermes

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__)

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.
 

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.
 

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):

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

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)