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: transaction_id = None else: transaction_id = transaction.id read_options = helpers.get_read_options(self._eventual, transaction_id) partition_id = entity_pb2.PartitionId( project_id=self._query.project, namespace_id=self._query.namespace) kwargs = {} if self._retry is not None: kwargs["retry"] = self._retry if self._timeout is not None: kwargs["timeout"] = self._timeout response_pb = self.client._datastore_api.run_query( request={ "project_id": self._query.project, "partition_id": partition_id, "read_options": read_options, "query": query_pb, }, **kwargs, ) while (response_pb.batch.more_results == _NOT_FINISHED and response_pb.batch.skipped_results < query_pb.offset): # We haven't finished processing. A likely reason is we haven't # skipped all of the results yet. Don't return any results. # Instead, rerun query, adjusting offsets. Datastore doesn't process # more than 1000 skipped results in a query. query_pb.start_cursor = response_pb.batch.skipped_cursor query_pb.offset -= response_pb.batch.skipped_results response_pb = self.client._datastore_api.run_query( request={ "project_id": self._query.project, "partition_id": partition_id, "read_options": read_options, "query": query_pb, }, **kwargs, ) entity_pbs = self._process_query_results(response_pb) return page_iterator.Page(self, entity_pbs, self.item_to_value)
def _next_page_helper(self, txn_id=None, retry=None, timeout=None): from google.api_core import page_iterator 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 from google.cloud.datastore.query import Query more_enum = query_pb2.QueryResultBatch.MoreResultsType.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) kwargs = {} if retry is not None: kwargs["retry"] = retry if timeout is not None: kwargs["timeout"] = timeout iterator = self._make_one(query, client, **kwargs) page = iterator._next_page() self.assertIsInstance(page, page_iterator.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( request={ "project_id": project, "partition_id": partition_id, "read_options": read_options, "query": empty_query, }, **kwargs, )
def test__next_page_w_skipped_lt_offset(self): from google.api_core import page_iterator 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 from google.cloud.datastore.query import Query project = "prujekt" skipped_1 = 100 skipped_cursor_1 = b"DEADBEEF" skipped_2 = 50 skipped_cursor_2 = b"FACEDACE" more_enum = query_pb2.QueryResultBatch.MoreResultsType.NOT_FINISHED result_1 = _make_query_response([], b"", more_enum, skipped_1) result_1.batch.skipped_cursor = skipped_cursor_1 result_2 = _make_query_response([], b"", more_enum, skipped_2) result_2.batch.skipped_cursor = skipped_cursor_2 ds_api = _make_datastore_api(result_1, result_2) client = _Client(project, datastore_api=ds_api) query = Query(client) offset = 150 iterator = self._make_one(query, client, offset=offset) page = iterator._next_page() self.assertIsInstance(page, page_iterator.Page) self.assertIs(page._parent, iterator) partition_id = entity_pb2.PartitionId(project_id=project) read_options = datastore_pb2.ReadOptions() query_1 = query_pb2.Query(offset=offset) query_2 = query_pb2.Query(start_cursor=skipped_cursor_1, offset=(offset - skipped_1)) expected_calls = [ mock.call( request={ "project_id": project, "partition_id": partition_id, "read_options": read_options, "query": query, }) for query in [query_1, query_2] ] self.assertEqual(ds_api.run_query.call_args_list, expected_calls)
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)