Skip to content

Commit

Permalink
Making datastore Connection.run_query() return low-level protobuf.
Browse files Browse the repository at this point in the history
Towards googleapis#2746 (as is googleapis#3064). This approach is to slowly transition
from our current approach to use the GAPIC generated surface.

These unit tests weren't so bad to update. I did "editorialize"
as I went, downgrading constants to just variables, renamining
for PEP8, etc.
  • Loading branch information
dhermes committed Feb 24, 2017
1 parent cd101b2 commit 222d08c
Show file tree
Hide file tree
Showing 4 changed files with 192 additions and 164 deletions.
14 changes: 3 additions & 11 deletions datastore/google/cloud/datastore/_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -379,10 +379,8 @@ def run_query(self, project, query_pb, namespace=None,
the given transaction. Incompatible with
``eventual==True``.
:rtype: tuple
:returns: Four-tuple containing the entities returned,
the end cursor of the query, a ``more_results``
enum and a count of the number of skipped results.
:rtype: :class:`.datastore_pb2.RunQueryResponse`
:returns: The protobuf response from a ``runQuery`` request.
"""
request = _datastore_pb2.RunQueryRequest()
_set_read_options(request, eventual, transaction_id)
Expand All @@ -391,13 +389,7 @@ def run_query(self, project, query_pb, namespace=None,
request.partition_id.namespace_id = namespace

request.query.CopyFrom(query_pb)
response = self._datastore_api.run_query(project, request)
return (
[e.entity for e in response.batch.entity_results],
response.batch.end_cursor, # Assume response always has cursor.
response.batch.more_results,
response.batch.skipped_results,
)
return self._datastore_api.run_query(project, request)

def begin_transaction(self, project):
"""Begin a transaction.
Expand Down
34 changes: 12 additions & 22 deletions datastore/google/cloud/datastore/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -441,43 +441,33 @@ def _build_protobuf(self):

return pb

def _process_query_results(self, entity_pbs, cursor_as_bytes,
more_results_enum, skipped_results):
def _process_query_results(self, response_pb):
"""Process the response from a datastore query.
:type entity_pbs: iterable
:param entity_pbs: The entities returned in the current page.
:type cursor_as_bytes: bytes
:param cursor_as_bytes: The end cursor of the query.
:type more_results_enum:
:class:`.query_pb2.QueryResultBatch.MoreResultsType`
:param more_results_enum: Enum indicating if there are more results.
:type skipped_results: int
:param skipped_results: The number of skipped results.
:type response_pb: :class:`.datastore_pb2.RunQueryResponse`
:param response_pb: The protobuf response from a ``runQuery`` request.
:rtype: iterable
:returns: The next page of entity results.
:raises ValueError: If ``more_results`` is an unexpected value.
"""
self._skipped_results = skipped_results
self._skipped_results = response_pb.batch.skipped_results

if cursor_as_bytes == b'': # Empty-value for bytes.
if response_pb.batch.end_cursor == b'': # Empty-value for bytes.
self.next_page_token = None
else:
self.next_page_token = base64.urlsafe_b64encode(cursor_as_bytes)
self.next_page_token = base64.urlsafe_b64encode(
response_pb.batch.end_cursor)
self._end_cursor = None

if more_results_enum == _NOT_FINISHED:
if response_pb.batch.more_results == _NOT_FINISHED:
self._more_results = True
elif more_results_enum in _FINISHED:
elif response_pb.batch.more_results in _FINISHED:
self._more_results = False
else:
raise ValueError('Unexpected value returned for `more_results`.')

return entity_pbs
return [result.entity for result in response_pb.batch.entity_results]

def _next_page(self):
"""Get the next page in the iterator.
Expand All @@ -492,13 +482,13 @@ def _next_page(self):
pb = self._build_protobuf()
transaction = self.client.current_transaction

query_results = self.client._connection.run_query(
response_pb = self.client._connection.run_query(
query_pb=pb,
project=self._query.project,
namespace=self._query.namespace,
transaction_id=transaction and transaction.id,
)
entity_pbs = self._process_query_results(*query_results)
entity_pbs = self._process_query_results(response_pb)
return Page(self, entity_pbs, self._item_to_value)


Expand Down
Loading

0 comments on commit 222d08c

Please sign in to comment.