Merge pull request #2097 from dhermes/datastore-prep-for-grpc

Factor out datastore API surface.

Author: dhermes

gcloud/datastore/connection.py

@@ -16,53 +16,27 @@
 
 import os
 
-from gcloud import connection
+from google.rpc import status_pb2
+
+from gcloud import connection as connection_module
 from gcloud.environment_vars import GCD_HOST
 from gcloud.exceptions import make_exception
 from gcloud.datastore._generated import datastore_pb2 as _datastore_pb2
-from google.rpc import status_pb2
-
 
-class Connection(connection.Connection):
- """A connection to the Google Cloud Datastore via the Protobuf API.
-
- This class should understand only the basic types (and protobufs)
- in method arguments, however should be capable of returning advanced types.
 
- :type credentials: :class:`oauth2client.client.OAuth2Credentials`
- :param credentials: The OAuth2 Credentials to use for this connection.
+class _DatastoreAPIOverHttp(object):
+ """Helper mapping datastore API methods.
 
- :type http: :class:`httplib2.Http` or class that defines ``request()``.
- :param http: An optional HTTP object to make requests.
+ Methods make bare API requests without any helpers for constructing
+ the requests or parsing the responses.
 
- :type api_base_url: string
- :param api_base_url: The base of the API call URL. Defaults to
-  :attr:`API_BASE_URL`.
+ :type connection: :class:`gcloud.datastore.connection.Connection`
+ :param connection: A connection object that contains helpful
+ information for making requests.
  """
 
- API_BASE_URL = 'https://datastore.googleapis.com'
- """The base of the API call URL."""
-
- API_VERSION = 'v1'
- """The version of the API, used in building the API call's URL."""
-
- API_URL_TEMPLATE = ('{api_base}/{api_version}/projects'
- '/{project}:{method}')
- """A template for the URL of a particular API call."""
-
- SCOPE = ('https://www.googleapis.com/auth/datastore',)
- """The scopes required for authenticating as a Cloud Datastore consumer."""
-
- def __init__(self, credentials=None, http=None, api_base_url=None):
- super(Connection, self).__init__(credentials=credentials, http=http)
- if api_base_url is None:
- try:
- # gcd.sh has /datastore/ in the path still since it supports
- # v1beta2 and v1beta3 simultaneously.
- api_base_url = '%s/datastore' % (os.environ[GCD_HOST],)
- except KeyError:
- api_base_url = self.__class__.API_BASE_URL
- self.api_base_url = api_base_url
+ def __init__(self, connection):
+ self.connection = connection
 
  def _request(self, project, method, data):
  """Make a request over the Http transport to the Cloud Datastore API.
@@ -86,10 +60,10 @@ def _request(self, project, method, data):
  headers = {
  'Content-Type': 'application/x-protobuf',
  'Content-Length': str(len(data)),
- 'User-Agent': self.USER_AGENT,
+ 'User-Agent': self.connection.USER_AGENT,
  }
- headers, content = self.http.request(
- uri=self.build_api_url(project=project, method=method),
+ headers, content = self.connection.http.request(
+ uri=self.connection.build_api_url(project=project, method=method),
  method='POST', headers=headers, body=data)
 
  status = headers['status']
@@ -124,6 +98,146 @@ def _rpc(self, project, method, request_pb, response_pb_cls):
  data=request_pb.SerializeToString())
  return response_pb_cls.FromString(response)
 
+ def lookup(self, project, request_pb):
+ """Perform a ``lookup`` request.
+
+ :type project: string
+ :param project: The project to connect to. This is
+ usually your project name in the cloud console.
+
+ :type request_pb: :class:`._generated.datastore_pb2.LookupRequest`
+ :param request_pb: The request protobuf object.
+
+ :rtype: :class:`._generated.datastore_pb2.LookupResponse`
+ :returns: The returned protobuf response object.
+ """
+ return self._rpc(project, 'lookup', request_pb,
+ _datastore_pb2.LookupResponse)
+
+ def run_query(self, project, request_pb):
+ """Perform a ``runQuery`` request.
+
+ :type project: string
+ :param project: The project to connect to. This is
+ usually your project name in the cloud console.
+
+ :type request_pb: :class:`._generated.datastore_pb2.RunQueryRequest`
+ :param request_pb: The request protobuf object.
+
+ :rtype: :class:`._generated.datastore_pb2.RunQueryResponse`
+ :returns: The returned protobuf response object.
+ """
+ return self._rpc(project, 'runQuery', request_pb,
+ _datastore_pb2.RunQueryResponse)
+
+ def begin_transaction(self, project, request_pb):
+ """Perform a ``beginTransaction`` request.
+
+ :type project: string
+ :param project: The project to connect to. This is
+ usually your project name in the cloud console.
+
+ :type request_pb:
+ :class:`._generated.datastore_pb2.BeginTransactionRequest`
+ :param request_pb: The request protobuf object.
+
+ :rtype: :class:`._generated.datastore_pb2.BeginTransactionResponse`
+ :returns: The returned protobuf response object.
+ """
+ return self._rpc(project, 'beginTransaction', request_pb,
+ _datastore_pb2.BeginTransactionResponse)
+
+ def commit(self, project, request_pb):
+ """Perform a ``commit`` request.
+
+ :type project: string
+ :param project: The project to connect to. This is
+ usually your project name in the cloud console.
+
+ :type request_pb: :class:`._generated.datastore_pb2.CommitRequest`
+ :param request_pb: The request protobuf object.
+
+ :rtype: :class:`._generated.datastore_pb2.CommitResponse`
+ :returns: The returned protobuf response object.
+ """
+ return self._rpc(project, 'commit', request_pb,
+ _datastore_pb2.CommitResponse)
+
+ def rollback(self, project, request_pb):
+ """Perform a ``rollback`` request.
+
+ :type project: string
+ :param project: The project to connect to. This is
+ usually your project name in the cloud console.
+
+ :type request_pb: :class:`._generated.datastore_pb2.RollbackRequest`
+ :param request_pb: The request protobuf object.
+
+ :rtype: :class:`._generated.datastore_pb2.RollbackResponse`
+ :returns: The returned protobuf response object.
+ """
+ return self._rpc(project, 'rollback', request_pb,
+ _datastore_pb2.RollbackResponse)
+
+ def allocate_ids(self, project, request_pb):
+ """Perform an ``allocateIds`` request.
+
+ :type project: string
+ :param project: The project to connect to. This is
+ usually your project name in the cloud console.
+
+ :type request_pb: :class:`._generated.datastore_pb2.AllocateIdsRequest`
+ :param request_pb: The request protobuf object.
+
+ :rtype: :class:`._generated.datastore_pb2.AllocateIdsResponse`
+ :returns: The returned protobuf response object.
+ """
+ return self._rpc(project, 'allocateIds', request_pb,
+ _datastore_pb2.AllocateIdsResponse)
+
+
+class Connection(connection_module.Connection):
+ """A connection to the Google Cloud Datastore via the Protobuf API.
+
+ This class should understand only the basic types (and protobufs)
+ in method arguments, however should be capable of returning advanced types.
+
+ :type credentials: :class:`oauth2client.client.OAuth2Credentials`
+ :param credentials: The OAuth2 Credentials to use for this connection.
+
+ :type http: :class:`httplib2.Http` or class that defines ``request()``.
+ :param http: An optional HTTP object to make requests.
+
+ :type api_base_url: string
+ :param api_base_url: The base of the API call URL. Defaults to
+ :attr:`API_BASE_URL`.
+ """
+
+ API_BASE_URL = 'https://datastore.googleapis.com'
+ """The base of the API call URL."""
+
+ API_VERSION = 'v1'
+ """The version of the API, used in building the API call's URL."""
+
+ API_URL_TEMPLATE = ('{api_base}/{api_version}/projects'
+ '/{project}:{method}')
+ """A template for the URL of a particular API call."""
+
+ SCOPE = ('https://www.googleapis.com/auth/datastore',)
+ """The scopes required for authenticating as a Cloud Datastore consumer."""
+
+ def __init__(self, credentials=None, http=None, api_base_url=None):
+ super(Connection, self).__init__(credentials=credentials, http=http)
+ if api_base_url is None:
+ try:
+ # gcd.sh has /datastore/ in the path still since it supports
+ # v1beta2 and v1beta3 simultaneously.
+ api_base_url = '%s/datastore' % (os.environ[GCD_HOST],)
+ except KeyError:
+ api_base_url = self.__class__.API_BASE_URL
+ self.api_base_url = api_base_url
+ self._datastore_api = _DatastoreAPIOverHttp(self)
+
  def build_api_url(self, project, method, base_url=None,
  api_version=None):
  """Construct the URL for a particular API call.
@@ -205,8 +319,7 @@ def lookup(self, project, key_pbs,
  _set_read_options(lookup_request, eventual, transaction_id)
  _add_keys_to_request(lookup_request.keys, key_pbs)
 
- lookup_response = self._rpc(project, 'lookup', lookup_request,
- _datastore_pb2.LookupResponse)
+ lookup_response = self._datastore_api.lookup(project, lookup_request)
 
  results = [result.entity for result in lookup_response.found]
  missing = [result.entity for result in lookup_response.missing]
@@ -260,8 +373,7 @@ def run_query(self, project, query_pb, namespace=None,
  request.partition_id.namespace_id = namespace
 
  request.query.CopyFrom(query_pb)
- response = self._rpc(project, 'runQuery', request,
- _datastore_pb2.RunQueryResponse)
+ response = self._datastore_api.run_query(project, request)
  return (
  [e.entity for e in response.batch.entity_results],
  response.batch.end_cursor, # Assume response always has cursor.
@@ -281,8 +393,7 @@ def begin_transaction(self, project):
  :returns: The serialized transaction that was begun.
  """
  request = _datastore_pb2.BeginTransactionRequest()
- response = self._rpc(project, 'beginTransaction', request,
- _datastore_pb2.BeginTransactionResponse)
+ response = self._datastore_api.begin_transaction(project, request)
  return response.transaction
 
  def commit(self, project, request, transaction_id):
@@ -316,8 +427,7 @@ def commit(self, project, request, transaction_id):
  else:
  request.mode = _datastore_pb2.CommitRequest.NON_TRANSACTIONAL
 
- response = self._rpc(project, 'commit', request,
- _datastore_pb2.CommitResponse)
+ response = self._datastore_api.commit(project, request)
  return _parse_commit_response(response)
 
  def rollback(self, project, transaction_id):
@@ -335,8 +445,7 @@ def rollback(self, project, transaction_id):
  request = _datastore_pb2.RollbackRequest()
  request.transaction = transaction_id
  # Nothing to do with this response, so just execute the method.
- self._rpc(project, 'rollback', request,
- _datastore_pb2.RollbackResponse)
+ self._datastore_api.rollback(project, request)
 
  def allocate_ids(self, project, key_pbs):
  """Obtain backend-generated IDs for a set of keys.
@@ -356,8 +465,7 @@ def allocate_ids(self, project, key_pbs):
  request = _datastore_pb2.AllocateIdsRequest()
  _add_keys_to_request(request.keys, key_pbs)
  # Nothing to do with this response, so just execute the method.
- response = self._rpc(project, 'allocateIds', request,
- _datastore_pb2.AllocateIdsResponse)
+ response = self._datastore_api.allocate_ids(project, request)
  return list(response.keys)