Add basic helpers needed for GAPIC client in datastore.

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)