def _next_page_helper(self, txn_id=None):
from google.cloud.iterator import Page
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
from google.cloud.datastore.query import Query
more_enum = query_pb2.QueryResultBatch.NOT_FINISHED
result = _make_query_response([], b'', more_enum, 0)
project = 'prujekt'
ds_api = _make_datastore_api(result)
if txn_id is None:
client = _Client(project, datastore_api=ds_api)
else:
transaction = mock.Mock(id=txn_id, spec=['id'])
client = _Client(
project, datastore_api=ds_api, transaction=transaction)
query = Query(client)
iterator = self._make_one(query, client)
page = iterator._next_page()
self.assertIsInstance(page, Page)
self.assertIs(page._parent, iterator)
partition_id = entity_pb2.PartitionId(project_id=project)
if txn_id is None:
read_options = datastore_pb2.ReadOptions()
else:
read_options = datastore_pb2.ReadOptions(transaction=txn_id)
empty_query = query_pb2.Query()
ds_api.run_query.assert_called_once_with(
project, partition_id, read_options, query=empty_query)
def _next_page(self):
"""Get the next page in the iterator.
:rtype: :class:`~google.cloud.iterator.Page`
:returns: The next page in the iterator (or :data:`None` if
there are no pages left).
"""
if not self._more_results:
return None
query_pb = self._build_protobuf()
transaction = self.client.current_transaction
if transaction is None:
read_options = _datastore_pb2.ReadOptions()
else:
read_options = _datastore_pb2.ReadOptions(
transaction=transaction.id)
partition_id = _entity_pb2.PartitionId(
project_id=self._query.project, namespace_id=self._query.namespace)
response_pb = self.client._datastore_api.run_query(
self._query.project,
partition_id,
read_options,
query=query_pb,
)
entity_pbs = self._process_query_results(response_pb)
return page_iterator.Page(self, entity_pbs, self._item_to_value)
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 make_partition(project, namespace):
"""Make a PartitionId for the given project and namespace."""
partition = entity_pb2.PartitionId()
partition.project_id = project
if namespace is not None:
partition.namespace_id = namespace
return partition
def test_run_query_exception(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 exception response
grpc_stub.RunQuery.side_effect = CustomException()
self.assertRaises(errors.GaxError, client.run_query, project_id,
partition_id)
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 run_data_migration():
request_data = json.loads(request.get_data())
# Required fields
fields = [
'name',
'function_kwargs',
'user'
]
# Some basic validation
for f in fields:
if f not in request_data:
resp_data = json.dumps(
{
'error': 'The ' + f + ' field is required.'
}
)
resp = Response(resp_data, status=400, mimetype='application/json')
return resp
if request_data['name'] not in migration.choices:
resp_data = json.dumps(
{
'error': 'The migration name is not valid.'
}
)
resp = Response(resp_data, status=400, mimetype='application/json')
return resp
migration_name = request_data['name']
function_kwargs = request_data['function_kwargs'] or {}
user = request_data['user']
function_kwargs.update({'name': migration_name})
# Create a MigrationHistory entity to keep track of the migration status
# set the project
project = PROJECT or 'meridianedit-staging'
# Create entity key
partition_id = entity_pb2.PartitionId(project_id=project, namespace_id="")
migration_history_obj_id = datetime.now().strftime("%Y%m%d%H%M%S")
path_element = entity_pb2.Key.PathElement(kind="MigrationHistory", name=migration_history_obj_id)
key = entity_pb2.Key(partition_id=partition_id, path=[path_element])
# Create entity and give it properties
entity = entity_pb2.Entity(key=key)
property_dict = {
'name': migration_name,
'function_kwargs': json.dumps(function_kwargs),
'started_by': user,
'status': 'running',
'created': datetime.now()
}
datastore_helper.add_properties(entity, property_dict)
# Add entity to datastore
mutations = [Mutation(insert=entity)]
client = apache_helper.get_datastore(project)
throttler = AdaptiveThrottler(window_ms=120000, bucket_ms=1000, overload_ratio=1.25)
apache_helper.write_mutations(client, project, mutations, throttler, rpc_stats_callback)
# Call the migration with any given function kwargs
migration_kwargs = {
'migration_history_obj': migration_history_obj_id,
}
migration_kwargs.update(function_kwargs)
# Run the migration in a celery task worker to prevent it timing
# out this connection. Also monitor the task so we can update
# migration status.
run_dataflow_migration.delay(pickle.dumps(entity), **migration_kwargs)
resp_data = {
'migration_history_obj_id': migration_history_obj_id
}
# A default 500 error message is returned if any of this breaks
return Response(json.dumps(resp_data), status=200, mimetype='application/json')
