@@ -441,43 +441,33 @@ def _build_protobuf(self):
441441
442442 return pb
443443
444- def _process_query_results (self , entity_pbs , cursor_as_bytes ,
445- more_results_enum , skipped_results ):
444+ def _process_query_results (self , response_pb ):
446445 """Process the response from a datastore query.
447446
448- :type entity_pbs: iterable
449- :param entity_pbs: The entities returned in the current page.
450-
451- :type cursor_as_bytes: bytes
452- :param cursor_as_bytes: The end cursor of the query.
453-
454- :type more_results_enum:
455- :class:`.query_pb2.QueryResultBatch.MoreResultsType`
456- :param more_results_enum: Enum indicating if there are more results.
457-
458- :type skipped_results: int
459- :param skipped_results: The number of skipped results.
447+ :type response_pb: :class:`.datastore_pb2.RunQueryResponse`
448+ :param response_pb: The protobuf response from a ``runQuery`` request.
460449
461450 :rtype: iterable
462451 :returns: The next page of entity results.
463452 :raises ValueError: If ``more_results`` is an unexpected value.
464453 """
465- self ._skipped_results = skipped_results
454+ self ._skipped_results = response_pb . batch . skipped_results
466455
467- if cursor_as_bytes == b'' : # Empty-value for bytes.
456+ if response_pb . batch . end_cursor == b'' : # Empty-value for bytes.
468457 self .next_page_token = None
469458 else :
470- self .next_page_token = base64 .urlsafe_b64encode (cursor_as_bytes )
459+ self .next_page_token = base64 .urlsafe_b64encode (
460+ response_pb .batch .end_cursor )
471461 self ._end_cursor = None
472462
473- if more_results_enum == _NOT_FINISHED :
463+ if response_pb . batch . more_results == _NOT_FINISHED :
474464 self ._more_results = True
475- elif more_results_enum in _FINISHED :
465+ elif response_pb . batch . more_results in _FINISHED :
476466 self ._more_results = False
477467 else :
478468 raise ValueError ('Unexpected value returned for `more_results`.' )
479469
480- return entity_pbs
470+ return [ result . entity for result in response_pb . batch . entity_results ]
481471
482472 def _next_page (self ):
483473 """Get the next page in the iterator.
@@ -492,13 +482,13 @@ def _next_page(self):
492482 pb = self ._build_protobuf ()
493483 transaction = self .client .current_transaction
494484
495- query_results = self .client ._connection .run_query (
485+ response_pb = self .client ._connection .run_query (
496486 query_pb = pb ,
497487 project = self ._query .project ,
498488 namespace = self ._query .namespace ,
499489 transaction_id = transaction and transaction .id ,
500490 )
501- entity_pbs = self ._process_query_results (* query_results )
491+ entity_pbs = self ._process_query_results (response_pb )
502492 return Page (self , entity_pbs , self ._item_to_value )
503493
504494
0 commit comments