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')