def test_run_query(self, mock_create_stub):
# Mock gRPC layer
grpc_stub = mock.Mock()
mock_create_stub.return_value = grpc_stub
client = datastore_client.DatastoreClient()
# Mock request
project_id = 'projectId-1969970175'
partition_id = entity_pb2.PartitionId()
# Mock response
expected_response = datastore_pb2.RunQueryResponse()
grpc_stub.RunQuery.return_value = expected_response
response = client.run_query(project_id, partition_id)
self.assertEqual(expected_response, response)
grpc_stub.RunQuery.assert_called_once()
args, kwargs = grpc_stub.RunQuery.call_args
self.assertEqual(len(args), 2)
self.assertEqual(len(kwargs), 1)
self.assertIn('metadata', kwargs)
actual_request = args[0]
expected_request = datastore_pb2.RunQueryRequest(
project_id=project_id, partition_id=partition_id)
self.assertEqual(expected_request, actual_request)
def test_run_query_w_namespace_nonempty_result(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.proto.datastore.v1 import entity_pb2
from google.cloud.proto.datastore.v1 import query_pb2
project = 'PROJECT'
kind = 'Kind'
entity = entity_pb2.Entity()
q_pb = self._make_query_pb(kind)
rsp_pb = datastore_pb2.RunQueryResponse()
rsp_pb.batch.entity_results.add(entity=entity)
rsp_pb.batch.entity_result_type = query_pb2.EntityResult.FULL
rsp_pb.batch.more_results = query_pb2.QueryResultBatch.NO_MORE_RESULTS
# Create mock HTTP and client with response.
http = Http({'status': '200'}, rsp_pb.SerializeToString())
client = mock.Mock(_http=http,
_base_url='test.invalid',
spec=['_http', '_base_url'])
# Make request.
conn = self._make_one(client)
namespace = 'NS'
response = conn.run_query(project, q_pb, namespace=namespace)
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb)
cw = http._called_with
uri = _build_expected_url(conn.api_base_url, project, 'runQuery')
_verify_protobuf_call(self, cw, uri)
request = datastore_pb2.RunQueryRequest()
request.ParseFromString(cw['body'])
self.assertEqual(request.partition_id.namespace_id, namespace)
self.assertEqual(request.query, q_pb)
def create_response(entities, end_cursor, finish):
"""Creates a query response for a given batch of scatter entities."""
resp = datastore_pb2.RunQueryResponse()
if finish:
resp.batch.more_results = query_pb2.QueryResultBatch.NO_MORE_RESULTS
else:
resp.batch.more_results = query_pb2.QueryResultBatch.NOT_FINISHED
resp.batch.end_cursor = end_cursor
for entity_result in entities:
resp.batch.entity_results.add().CopyFrom(entity_result)
return resp
def _make_query_response(entity_pbs, cursor_as_bytes, more_results_enum,
skipped_results):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.proto.datastore.v1 import query_pb2
return datastore_pb2.RunQueryResponse(batch=query_pb2.QueryResultBatch(
skipped_results=skipped_results,
end_cursor=cursor_as_bytes,
more_results=more_results_enum,
entity_results=[
query_pb2.EntityResult(entity=entity) for entity in entity_pbs
],
), )
def test_run_query_w_namespace_nonempty_result(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.proto.datastore.v1 import entity_pb2
from google.cloud.proto.datastore.v1 import query_pb2
project = 'PROJECT'
kind = 'Kind'
namespace = 'NS'
query_pb = self._make_query_pb(kind)
partition_id = entity_pb2.PartitionId(project_id=project,
namespace_id=namespace)
read_options = datastore_pb2.ReadOptions()
rsp_pb = datastore_pb2.RunQueryResponse(
batch=query_pb2.QueryResultBatch(
entity_result_type=query_pb2.EntityResult.FULL,
entity_results=[
query_pb2.EntityResult(entity=entity_pb2.Entity()),
],
more_results=query_pb2.QueryResultBatch.NO_MORE_RESULTS,
))
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb.SerializeToString())])
client = mock.Mock(_http=http,
_base_url='test.invalid',
spec=['_http', '_base_url'])
# Make request.
ds_api = self._make_one(client)
response = ds_api.run_query(project,
partition_id,
read_options,
query=query_pb)
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb)
uri = _build_expected_url(client._base_url, project, 'runQuery')
request = _verify_protobuf_call(http, uri,
datastore_pb2.RunQueryRequest())
self.assertEqual(request.partition_id, partition_id)
self.assertEqual(request.query, query_pb)
def test_run_query_wo_eventual_w_transaction(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.proto.datastore.v1 import entity_pb2
from google.cloud.proto.datastore.v1 import query_pb2
project = 'PROJECT'
kind = 'Nonesuch'
cursor = b'\x00'
transaction = b'TRANSACTION'
query_pb = self._make_query_pb(kind)
partition_id = entity_pb2.PartitionId(project_id=project)
read_options = datastore_pb2.ReadOptions(transaction=transaction)
rsp_pb = datastore_pb2.RunQueryResponse(
batch=query_pb2.QueryResultBatch(
entity_result_type=query_pb2.EntityResult.FULL,
end_cursor=cursor,
more_results=query_pb2.QueryResultBatch.NO_MORE_RESULTS,
))
# Create mock HTTP and client with response.
http = Http({'status': '200'}, rsp_pb.SerializeToString())
client = mock.Mock(_http=http,
_base_url='test.invalid',
spec=['_http', '_base_url'])
# Make request.
ds_api = self._make_one(client)
response = ds_api.run_query(project,
partition_id,
read_options,
query=query_pb)
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb)
uri = _build_expected_url(client._base_url, project, 'runQuery')
cw = http._called_with
_verify_protobuf_call(self, cw, uri)
request = datastore_pb2.RunQueryRequest()
request.ParseFromString(cw['body'])
self.assertEqual(request.partition_id, partition_id)
self.assertEqual(request.query, query_pb)
self.assertEqual(request.read_options, read_options)
def test_run_query_wo_eventual_w_transaction(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.proto.datastore.v1 import query_pb2
project = 'PROJECT'
kind = 'Nonesuch'
cursor = b'\x00'
transaction = b'TRANSACTION'
q_pb = self._make_query_pb(kind)
rsp_pb = datastore_pb2.RunQueryResponse()
rsp_pb.batch.end_cursor = cursor
no_more = query_pb2.QueryResultBatch.NO_MORE_RESULTS
rsp_pb.batch.more_results = no_more
rsp_pb.batch.entity_result_type = query_pb2.EntityResult.FULL
# Create mock HTTP and client with response.
http = Http({'status': '200'}, rsp_pb.SerializeToString())
client = mock.Mock(_http=http, spec=['_http'])
# Make request.
conn = self._make_one(client)
response = conn.run_query(project, q_pb, transaction_id=transaction)
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb)
uri = '/'.join([
conn.api_base_url,
conn.API_VERSION,
'projects',
project + ':runQuery',
])
cw = http._called_with
self._verify_protobuf_call(cw, uri, conn)
request = datastore_pb2.RunQueryRequest()
request.ParseFromString(cw['body'])
self.assertEqual(request.partition_id.namespace_id, '')
self.assertEqual(request.query, q_pb)
self.assertEqual(
request.read_options.read_consistency,
datastore_pb2.ReadOptions.READ_CONSISTENCY_UNSPECIFIED)
self.assertEqual(request.read_options.transaction, transaction)
def test_run_query_w_eventual_and_transaction(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.proto.datastore.v1 import query_pb2
PROJECT = 'PROJECT'
KIND = 'Nonesuch'
CURSOR = b'\x00'
TRANSACTION = b'TRANSACTION'
q_pb = self._make_query_pb(KIND)
rsp_pb = datastore_pb2.RunQueryResponse()
rsp_pb.batch.end_cursor = CURSOR
no_more = query_pb2.QueryResultBatch.NO_MORE_RESULTS
rsp_pb.batch.more_results = no_more
rsp_pb.batch.entity_result_type = query_pb2.EntityResult.FULL
conn = self._make_one(object())
self.assertRaises(ValueError,
conn.run_query,
PROJECT,
q_pb,
eventual=True,
transaction_id=TRANSACTION)
def test_run_query_w_eventual_no_transaction(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.proto.datastore.v1 import query_pb2
project = 'PROJECT'
kind = 'Nonesuch'
cursor = b'\x00'
q_pb = self._make_query_pb(kind)
rsp_pb = datastore_pb2.RunQueryResponse()
rsp_pb.batch.end_cursor = cursor
no_more = query_pb2.QueryResultBatch.NO_MORE_RESULTS
rsp_pb.batch.more_results = no_more
rsp_pb.batch.entity_result_type = query_pb2.EntityResult.FULL
# Create mock HTTP and client with response.
http = Http({'status': '200'}, rsp_pb.SerializeToString())
client = mock.Mock(_http=http,
_base_url='test.invalid',
spec=['_http', '_base_url'])
# Make request.
conn = self._make_one(client)
response = conn.run_query(project, q_pb, eventual=True)
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb)
uri = _build_expected_url(conn.api_base_url, project, 'runQuery')
cw = http._called_with
_verify_protobuf_call(self, cw, uri)
request = datastore_pb2.RunQueryRequest()
request.ParseFromString(cw['body'])
self.assertEqual(request.partition_id.namespace_id, '')
self.assertEqual(request.query, q_pb)
self.assertEqual(request.read_options.read_consistency,
datastore_pb2.ReadOptions.EVENTUAL)
self.assertEqual(request.read_options.transaction, b'')
def make_stats_response(self, property_map):
resp = datastore_pb2.RunQueryResponse()
entity_result = resp.batch.entity_results.add()
datastore_helper.add_properties(entity_result.entity, property_map)
return resp
def transient_retriable_datastore_failure(self, req):
if self._transient_fail_count:
self._transient_fail_count -= 1
raise random.choice(self._retriable_errors)
else:
return datastore_pb2.RunQueryResponse()
def transient_retriable_datastore_failure(self, req):
if self._transient_fail_count:
self._transient_fail_count -= 1
raise RPCError("dummy", code_pb2.INTERNAL, "failed")
else:
return datastore_pb2.RunQueryResponse()
def transient_datastore_failure(self, req):
if self._transient_fail_count:
self._transient_fail_count -= 1
raise RPCError("dummy", 500, "failed")
else:
return datastore_pb2.RunQueryResponse()