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

Factor out datastore API surface.

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)