def _make_query_response(entity_pbs, cursor_as_bytes, more_results_enum,
skipped_results):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import query as 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.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore_v1.types import query as 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.ResultType.FULL,
entity_results=[
query_pb2.EntityResult(entity=entity_pb2.Entity())
],
more_results=query_pb2.QueryResultBatch.MoreResultsType.
NO_MORE_RESULTS,
))
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb._pb.SerializeToString())])
client_info = _make_client_info()
client = mock.Mock(
_http=http,
_base_url="test.invalid",
_client_info=client_info,
spec=["_http", "_base_url", "_client_info"],
)
# 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._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._pb)
self.assertEqual(request.query, query_pb._pb)
def _run_query_helper(
self,
read_consistency=None,
transaction=None,
namespace=None,
found=0,
retry=None,
timeout=None,
):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore_v1.types import query as query_pb2
project = "PROJECT"
kind = "Nonesuch"
query_pb = self._make_query_pb(kind)
partition_kw = {"project_id": project}
if namespace is not None:
partition_kw["namespace_id"] = namespace
partition_id = entity_pb2.PartitionId(**partition_kw)
options_kw = {}
if read_consistency is not None:
options_kw["read_consistency"] = read_consistency
if transaction is not None:
options_kw["transaction"] = transaction
read_options = datastore_pb2.ReadOptions(**options_kw)
cursor = b"\x00"
batch_kw = {
"entity_result_type":
query_pb2.EntityResult.ResultType.FULL,
"end_cursor":
cursor,
"more_results":
query_pb2.QueryResultBatch.MoreResultsType.NO_MORE_RESULTS,
}
if found:
batch_kw["entity_results"] = [
query_pb2.EntityResult(entity=entity_pb2.Entity())
] * found
rsp_pb = datastore_pb2.RunQueryResponse(
batch=query_pb2.QueryResultBatch(**batch_kw))
http = _make_requests_session(
[_make_response(content=rsp_pb._pb.SerializeToString())])
client_info = _make_client_info()
client = mock.Mock(
_http=http,
_base_url="test.invalid",
_client_info=client_info,
spec=["_http", "_base_url", "_client_info"],
)
ds_api = self._make_one(client)
request = {
"project_id": project,
"partition_id": partition_id,
"read_options": read_options,
"query": query_pb,
}
kwargs = _make_retry_timeout_kwargs(retry, timeout, http)
response = ds_api.run_query(request=request, **kwargs)
self.assertEqual(response, rsp_pb._pb)
uri = _build_expected_url(client._base_url, project, "runQuery")
request = _verify_protobuf_call(
http,
uri,
datastore_pb2.RunQueryRequest(),
retry=retry,
timeout=timeout,
)
self.assertEqual(request.partition_id, partition_id._pb)
self.assertEqual(request.query, query_pb._pb)
self.assertEqual(request.read_options, read_options._pb)