diff --git a/gcloud/datastore/connection.py b/gcloud/datastore/connection.py index e2788fee08bc..9fcacacfcd6f 100644 --- a/gcloud/datastore/connection.py +++ b/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) diff --git a/gcloud/datastore/test_connection.py b/gcloud/datastore/test_connection.py index 1c8be40a57f3..e4e3549cdfcf 100644 --- a/gcloud/datastore/test_connection.py +++ b/gcloud/datastore/test_connection.py @@ -15,6 +15,103 @@ import unittest +class Test_DatastoreAPIOverHttp(unittest.TestCase): + + def _getTargetClass(self): + from gcloud.datastore.connection import _DatastoreAPIOverHttp + return _DatastoreAPIOverHttp + + def _makeOne(self, *args, **kw): + return self._getTargetClass()(*args, **kw) + + def test__rpc(self): + from gcloud.datastore.connection import Connection + + class ReqPB(object): + + def SerializeToString(self): + return REQPB + + class RspPB(object): + + def __init__(self, pb): + self._pb = pb + + @classmethod + def FromString(cls, pb): + return cls(pb) + + REQPB = b'REQPB' + PROJECT = 'PROJECT' + METHOD = 'METHOD' + conn = Connection() + datastore_api = self._makeOne(conn) + URI = '/'.join([ + conn.api_base_url, + conn.API_VERSION, + 'projects', + PROJECT + ':' + METHOD, + ]) + http = conn._http = Http({'status': '200'}, 'CONTENT') + response = datastore_api._rpc(PROJECT, METHOD, ReqPB(), RspPB) + self.assertTrue(isinstance(response, RspPB)) + self.assertEqual(response._pb, 'CONTENT') + called_with = http._called_with + self.assertEqual(called_with['uri'], URI) + self.assertEqual(called_with['method'], 'POST') + self.assertEqual(called_with['headers']['Content-Type'], + 'application/x-protobuf') + self.assertEqual(called_with['headers']['User-Agent'], + conn.USER_AGENT) + self.assertEqual(called_with['body'], REQPB) + + def test__request_w_200(self): + from gcloud.datastore.connection import Connection + + PROJECT = 'PROJECT' + METHOD = 'METHOD' + DATA = b'DATA' + conn = Connection() + datastore_api = self._makeOne(conn) + URI = '/'.join([ + conn.api_base_url, + conn.API_VERSION, + 'projects', + PROJECT + ':' + METHOD, + ]) + http = conn._http = Http({'status': '200'}, 'CONTENT') + self.assertEqual(datastore_api._request(PROJECT, METHOD, DATA), + 'CONTENT') + called_with = http._called_with + self.assertEqual(called_with['uri'], URI) + self.assertEqual(called_with['method'], 'POST') + self.assertEqual(called_with['headers']['Content-Type'], + 'application/x-protobuf') + self.assertEqual(called_with['headers']['User-Agent'], + conn.USER_AGENT) + self.assertEqual(called_with['body'], DATA) + + def test__request_not_200(self): + from gcloud.datastore.connection import Connection + from gcloud.exceptions import BadRequest + from google.rpc import status_pb2 + + error = status_pb2.Status() + error.message = 'Entity value is indexed.' + error.code = 9 # FAILED_PRECONDITION + + PROJECT = 'PROJECT' + METHOD = 'METHOD' + DATA = 'DATA' + conn = Connection() + datastore_api = self._makeOne(conn) + conn._http = Http({'status': '400'}, error.SerializeToString()) + with self.assertRaises(BadRequest) as e: + datastore_api._request(PROJECT, METHOD, DATA) + expected_message = '400 Entity value is indexed.' + self.assertEqual(str(e.exception), expected_message) + + class TestConnection(unittest.TestCase): def _getTargetClass(self): @@ -135,73 +232,6 @@ def create_scoped_required(self): self.assertTrue(conn.http is authorized) self.assertTrue(isinstance(creds._called_with, httplib2.Http)) - def test__request_w_200(self): - PROJECT = 'PROJECT' - METHOD = 'METHOD' - DATA = b'DATA' - conn = self._makeOne() - URI = '/'.join([ - conn.api_base_url, - conn.API_VERSION, - 'projects', - PROJECT + ':' + METHOD, - ]) - http = conn._http = Http({'status': '200'}, 'CONTENT') - self.assertEqual(conn._request(PROJECT, METHOD, DATA), 'CONTENT') - self._verifyProtobufCall(http._called_with, URI, conn) - self.assertEqual(http._called_with['body'], DATA) - - def test__request_not_200(self): - from gcloud.exceptions import BadRequest - from google.rpc import status_pb2 - - error = status_pb2.Status() - error.message = 'Entity value is indexed.' - error.code = 9 # FAILED_PRECONDITION - - PROJECT = 'PROJECT' - METHOD = 'METHOD' - DATA = 'DATA' - conn = self._makeOne() - conn._http = Http({'status': '400'}, error.SerializeToString()) - with self.assertRaises(BadRequest) as e: - conn._request(PROJECT, METHOD, DATA) - expected_message = '400 Entity value is indexed.' - self.assertEqual(str(e.exception), expected_message) - - def test__rpc(self): - - class ReqPB(object): - - def SerializeToString(self): - return REQPB - - class RspPB(object): - - def __init__(self, pb): - self._pb = pb - - @classmethod - def FromString(cls, pb): - return cls(pb) - - REQPB = b'REQPB' - PROJECT = 'PROJECT' - METHOD = 'METHOD' - conn = self._makeOne() - URI = '/'.join([ - conn.api_base_url, - conn.API_VERSION, - 'projects', - PROJECT + ':' + METHOD, - ]) - http = conn._http = Http({'status': '200'}, 'CONTENT') - response = conn._rpc(PROJECT, METHOD, ReqPB(), RspPB) - self.assertTrue(isinstance(response, RspPB)) - self.assertEqual(response._pb, 'CONTENT') - self._verifyProtobufCall(http._called_with, URI, conn) - self.assertEqual(http._called_with['body'], REQPB) - def test_build_api_url_w_default_base_version(self): PROJECT = 'PROJECT' METHOD = 'METHOD'