diff --git a/firestore/google/cloud/firestore_v1beta1/query.py b/firestore/google/cloud/firestore_v1beta1/query.py index 489cc1f82a05..ebf2b061945c 100644 --- a/firestore/google/cloud/firestore_v1beta1/query.py +++ b/firestore/google/cloud/firestore_v1beta1/query.py @@ -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): @@ -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( @@ -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): @@ -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) @@ -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 diff --git a/firestore/tests/unit/test_query.py b/firestore/tests/unit/test_query.py index 0b4e6e08e6fe..39fae3741113 100644 --- a/firestore/tests/unit/test_query.py +++ b/firestore/tests/unit/test_query.py @@ -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() @@ -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"]) @@ -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() @@ -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 @@ -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)