Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 9 additions & 33 deletions firestore/google/cloud/firestore_v1beta1/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,10 +63,6 @@
"come from fields set in ``order_by()``."
)
_MISMATCH_CURSOR_W_ORDER_BY = "The cursor {!r} does not match the order fields {!r}."
_EMPTY_DOC_TEMPLATE = (
"Unexpected server response. All responses other than the first must "
"contain a document. The response at index {} was\n{}."
)


class Query(object):
Expand Down Expand Up @@ -725,12 +721,6 @@ def get(self, transaction=None):
Yields:
~.firestore_v1beta1.document.DocumentSnapshot: The next
document that fulfills the query.

Raises:
ValueError: If the first response in the stream is empty, but
then more responses follow.
ValueError: If a response other than the first does not contain
a document.
"""
parent_path, expected_prefix = self._parent._parent_info()
response_iterator = self._client._firestore_api.run_query(
Expand All @@ -740,24 +730,11 @@ def get(self, transaction=None):
metadata=self._client._rpc_metadata,
)

empty_stream = False
for index, response_pb in enumerate(response_iterator):
if empty_stream:
raise ValueError(
"First response in stream was empty",
"Received second response",
response_pb,
)

snapshot, skipped_results = _query_response_to_snapshot(
response_pb, self._parent, expected_prefix
for response in response_iterator:
snapshot = _query_response_to_snapshot(
response, self._parent, expected_prefix
)
if snapshot is None:
if index != 0:
msg = _EMPTY_DOC_TEMPLATE.format(index, response_pb)
raise ValueError(msg)
empty_stream = skipped_results == 0
else:
if snapshot is not None:
yield snapshot

def on_snapshot(self, callback):
Expand Down Expand Up @@ -964,13 +941,12 @@ def _query_response_to_snapshot(response_pb, collection, expected_prefix):
directly from ``collection`` via :meth:`_parent_info`.

Returns:
Tuple[Optional[~.firestore.document.DocumentSnapshot], int]: A
snapshot of the data returned in the query and the number of skipped
results. If ``response_pb.document`` is not set, the snapshot will be
:data:`None`.
Optional[~.firestore.document.DocumentSnapshot]: A
snapshot of the data returned in the query. If ``response_pb.document``
is not set, the snapshot will be :data:`None`.
"""
if not response_pb.HasField("document"):
return None, response_pb.skipped_results
return None

document_id = _helpers.get_doc_id(response_pb.document, expected_prefix)
reference = collection.document(document_id)
Expand All @@ -983,4 +959,4 @@ def _query_response_to_snapshot(response_pb, collection, expected_prefix):
create_time=response_pb.document.create_time,
update_time=response_pb.document.update_time,
)
return snapshot, response_pb.skipped_results
return snapshot
33 changes: 9 additions & 24 deletions firestore/tests/unit/test_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -1137,13 +1137,7 @@ def test_get_second_response_in_empty_stream(self):

get_response = query.get()
self.assertIsInstance(get_response, types.GeneratorType)
with self.assertRaises(ValueError) as exc_info:
list(get_response)

exc_args = exc_info.exception.args
self.assertEqual(len(exc_args), 3)
self.assertIs(exc_args[2], empty_response2)
self.assertIsNot(empty_response1, empty_response2)
self.assertEqual(list(get_response), [])

# Verify the mock call.
parent_path, _ = parent._parent_info()
Expand Down Expand Up @@ -1193,8 +1187,6 @@ def test_get_with_skipped_results(self):
)

def test_get_empty_after_first_response(self):
from google.cloud.firestore_v1beta1.query import _EMPTY_DOC_TEMPLATE

# Create a minimal fake GAPIC.
firestore_api = mock.Mock(spec=["run_query"])

Expand All @@ -1217,13 +1209,11 @@ def test_get_empty_after_first_response(self):
query = self._make_one(parent)
get_response = query.get()
self.assertIsInstance(get_response, types.GeneratorType)
with self.assertRaises(ValueError) as exc_info:
list(get_response)

exc_args = exc_info.exception.args
self.assertEqual(len(exc_args), 1)
msg = _EMPTY_DOC_TEMPLATE.format(1, response_pb2)
self.assertEqual(exc_args[0], msg)
returned = list(get_response)
self.assertEqual(len(returned), 1)
snapshot = returned[0]
self.assertEqual(snapshot.reference._path, ("charles", "bark"))
self.assertEqual(snapshot.to_dict(), data)

# Verify the mock call.
parent_path, _ = parent._parent_info()
Expand Down Expand Up @@ -1468,16 +1458,14 @@ def _call_fut(response_pb, collection, expected_prefix):

def test_empty(self):
response_pb = _make_query_response()
snapshot, skipped_results = self._call_fut(response_pb, None, None)
snapshot = self._call_fut(response_pb, None, None)
self.assertIsNone(snapshot)
self.assertEqual(skipped_results, 0)

def test_after_offset(self):
skipped_results = 410
response_pb = _make_query_response(skipped_results=skipped_results)
snapshot, skipped_results = self._call_fut(response_pb, None, None)
snapshot = self._call_fut(response_pb, None, None)
self.assertIsNone(snapshot)
self.assertEqual(skipped_results, skipped_results)

def test_response(self):
from google.cloud.firestore_v1beta1.document import DocumentSnapshot
Expand All @@ -1492,10 +1480,7 @@ def test_response(self):
data = {"a": 901, "b": True}
response_pb = _make_query_response(name=name, data=data)

snapshot, skipped_results = self._call_fut(
response_pb, collection, expected_prefix
)
self.assertEqual(skipped_results, 0)
snapshot = self._call_fut(response_pb, collection, expected_prefix)
self.assertIsInstance(snapshot, DocumentSnapshot)
expected_path = collection._path + (doc_id,)
self.assertEqual(snapshot.reference._path, expected_path)
Expand Down