feat: use generator for stream results (#926)

* feat: use iterator for query results

* use iterator as suggested in comments

* remove unnecessary class

* more iterators

* undo adding await

* address comments

* undo bundle change

* undo bundle change

* cleanups and docstrings

* fix type hint

* unit tests

* lint

* skip tests with anext for python < 3.10

* lint

* address comments

* lint

* fix type hint

* type hint

* sys test debug

* sys test debug

* undo change for debug

* address comment

* system test debug

* undo system test debug code

Author: Linchin

google/cloud/firestore_v1/aggregation.py

@@ -30,8 +30,14 @@
     BaseAggregationQuery,
     _query_response_to_result,
 )
+from google.cloud.firestore_v1.base_document import DocumentSnapshot
+from google.cloud.firestore_v1.stream_generator import StreamGenerator
 
-from typing import Generator, Union, List, Any
+from typing import Any, Generator, List, Optional, TYPE_CHECKING, Union
+
+# Types needed only for Type Hints
+if TYPE_CHECKING:
+    from google.cloud.firestore_v1 import transaction  # pragma: NO COVER
 
 
 class AggregationQuery(BaseAggregationQuery):
@@ -99,36 +105,34 @@ def _retry_query_after_exception(self, exc, retry, transaction):
 
         return False
 
-    def stream(
+    def _make_stream(
         self,
-        transaction=None,
-        retry: Union[
-            retries.Retry, None, gapic_v1.method._MethodDefault
-        ] = gapic_v1.method.DEFAULT,
-        timeout: float | None = None,
+        transaction: Optional[transaction.Transaction] = None,
+        retry: Optional[retries.Retry] = gapic_v1.method.DEFAULT,
+        timeout: Optional[float] = None,
     ) -> Union[Generator[List[AggregationResult], Any, None]]:
-        """Runs the aggregation query.
+        """Internal method for stream(). Runs the aggregation query.
 
-        This sends a ``RunAggregationQuery`` RPC and then returns an iterator which
-        consumes each document returned in the stream of ``RunAggregationQueryResponse``
-        messages.
+        This sends a ``RunAggregationQuery`` RPC and then returns a generator
+        which consumes each document returned in the stream of
+        ``RunAggregationQueryResponse`` messages.
 
-        If a ``transaction`` is used and it already has write operations
-        added, this method cannot be used (i.e. read-after-write is not
-        allowed).
+        If a ``transaction`` is used and it already has write operations added,
+        this method cannot be used (i.e. read-after-write is not allowed).
 
         Args:
             transaction
                 (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
                 An existing transaction that this query will run in.
-            retry (google.api_core.retry.Retry): Designation of what errors, if any,
-                should be retried.  Defaults to a system-specified policy.
-            timeout (float): The timeout for this request.  Defaults to a
-                system-specified value.
+            retry (Optional[google.api_core.retry.Retry]): Designation of what
+                errors, if any, should be retried.  Defaults to a
+                system-specified policy.
+            timeout (Optional[float]): The timeout for this request.  Defaults
+                to a system-specified value.
 
         Yields:
             :class:`~google.cloud.firestore_v1.base_aggregation.AggregationResult`:
-            The result of aggregations of this query
+            The result of aggregations of this query.
         """
 
         response_iterator = self._get_stream_iterator(
@@ -154,3 +158,38 @@ def stream(
                 break
             result = _query_response_to_result(response)
             yield result
+
+    def stream(
+        self,
+        transaction: Optional["transaction.Transaction"] = None,
+        retry: Optional[retries.Retry] = gapic_v1.method.DEFAULT,
+        timeout: Optional[float] = None,
+    ) -> "StreamGenerator[DocumentSnapshot]":
+        """Runs the aggregation query.
+
+        This sends a ``RunAggregationQuery`` RPC and then returns a generator
+        which consumes each document returned in the stream of
+        ``RunAggregationQueryResponse`` messages.
+
+        If a ``transaction`` is used and it already has write operations added,
+        this method cannot be used (i.e. read-after-write is not allowed).
+
+        Args:
+            transaction
+                (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
+                An existing transaction that this query will run in.
+            retry (Optional[google.api_core.retry.Retry]): Designation of what
+                errors, if any, should be retried.  Defaults to a
+                system-specified policy.
+            timeout (Optinal[float]): The timeout for this request.  Defaults
+            to a system-specified value.
+
+        Returns:
+            `StreamGenerator[DocumentSnapshot]`: A generator of the query results.
+        """
+        inner_generator = self._make_stream(
+            transaction=transaction,
+            retry=retry,
+            timeout=timeout,
+        )
+        return StreamGenerator(inner_generator)

google/cloud/firestore_v1/async_aggregation.py

@@ -23,14 +23,18 @@
 from google.api_core import gapic_v1
 from google.api_core import retry_async as retries
 
-from typing import List, Union, AsyncGenerator
-
+from typing import AsyncGenerator, List, Optional, Union, TYPE_CHECKING
 
+from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator
 from google.cloud.firestore_v1.base_aggregation import (
     AggregationResult,
     _query_response_to_result,
     BaseAggregationQuery,
 )
+from google.cloud.firestore_v1 import transaction
+
+if TYPE_CHECKING:  # pragma: NO COVER
+    from google.cloud.firestore_v1.base_document import DocumentSnapshot
 
 
 class AsyncAggregationQuery(BaseAggregationQuery):
@@ -76,17 +80,15 @@ async def get(
         result = [aggregation async for aggregation in stream_result]
         return result  # type: ignore
 
-    async def stream(
+    async def _make_stream(
         self,
-        transaction=None,
-        retry: Union[
-            retries.AsyncRetry, None, gapic_v1.method._MethodDefault
-        ] = gapic_v1.method.DEFAULT,
-        timeout: float | None = None,
+        transaction: Optional[transaction.Transaction] = None,
+        retry: Optional[retries.AsyncRetry] = gapic_v1.method.DEFAULT,
+        timeout: Optional[float] = None,
     ) -> Union[AsyncGenerator[List[AggregationResult], None]]:
-        """Runs the aggregation query.
+        """Internal method for stream(). Runs the aggregation query.
 
-        This sends a ``RunAggregationQuery`` RPC and then returns an iterator which
+        This sends a ``RunAggregationQuery`` RPC and then returns a generator which
         consumes each document returned in the stream of ``RunAggregationQueryResponse``
         messages.
 
@@ -95,13 +97,14 @@ async def stream(
         allowed).
 
         Args:
-            transaction
-                (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
-                An existing transaction that this query will run in.
-            retry (google.api_core.retry.Retry): Designation of what errors, if any,
-                should be retried.  Defaults to a system-specified policy.
-            timeout (float): The timeout for this request.  Defaults to a
-                system-specified value.
+            transaction (Optional[:class:`~google.cloud.firestore_v1.transaction.\
+                Transaction`]):
+                An existing transaction that the query will run in.
+            retry (Optional[google.api_core.retry.Retry]): Designation of what
+                errors, if any, should be retried.  Defaults to a
+                system-specified policy.
+            timeout (Optional[float]): The timeout for this request. Defaults
+                to a system-specified value.
 
         Yields:
             :class:`~google.cloud.firestore_v1.base_aggregation.AggregationResult`:
@@ -122,3 +125,40 @@ async def stream(
         async for response in response_iterator:
             result = _query_response_to_result(response)
             yield result
+
+    def stream(
+        self,
+        transaction: Optional[transaction.Transaction] = None,
+        retry: Optional[retries.AsyncRetry] = gapic_v1.method.DEFAULT,
+        timeout: Optional[float] = None,
+    ) -> "AsyncStreamGenerator[DocumentSnapshot]":
+        """Runs the aggregation query.
+
+        This sends a ``RunAggregationQuery`` RPC and then returns a generator
+        which consumes each document returned in the stream of
+        ``RunAggregationQueryResponse`` messages.
+
+        If a ``transaction`` is used and it already has write operations added,
+        this method cannot be used (i.e. read-after-write is not allowed).
+
+        Args:
+            transaction (Optional[:class:`~google.cloud.firestore_v1.transaction.\
+                Transaction`]):
+                An existing transaction that the query will run in.
+            retry (Optional[google.api_core.retry.Retry]): Designation of what
+                errors, if any, should be retried.  Defaults to a
+                system-specified policy.
+            timeout (Optional[float]): The timeout for this request. Defaults
+                to a system-specified value.
+
+        Returns:
+            `AsyncStreamGenerator[DocumentSnapshot]`:
+                A generator of the query results.
+        """
+
+        inner_generator = self._make_stream(
+            transaction=transaction,
+            retry=retry,
+            timeout=timeout,
+        )
+        return AsyncStreamGenerator(inner_generator)

google/cloud/firestore_v1/async_collection.py

@@ -21,15 +21,20 @@
     BaseCollectionReference,
     _item_to_document_ref,
 )
-from google.cloud.firestore_v1 import async_query, async_document, async_aggregation
+from google.cloud.firestore_v1 import (
+    async_aggregation,
+    async_document,
+    async_query,
+    transaction,
+)
 
 from google.cloud.firestore_v1.document import DocumentReference
 
-from typing import AsyncIterator
-from typing import Any, AsyncGenerator, Tuple
+from typing import Any, AsyncGenerator, Optional, Tuple, TYPE_CHECKING
 
-# Types needed only for Type Hints
-from google.cloud.firestore_v1.transaction import Transaction
+if TYPE_CHECKING:  # pragma: NO COVER
+    from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator
+    from google.cloud.firestore_v1.base_document import DocumentSnapshot
 
 
 class AsyncCollectionReference(BaseCollectionReference[async_query.AsyncQuery]):
@@ -176,9 +181,9 @@ async def list_documents(
 
     async def get(
         self,
-        transaction: Transaction = None,
-        retry: retries.AsyncRetry = gapic_v1.method.DEFAULT,
-        timeout: float = None,
+        transaction: Optional[transaction.Transaction] = None,
+        retry: Optional[retries.AsyncRetry] = gapic_v1.method.DEFAULT,
+        timeout: Optional[float] = None,
     ) -> list:
         """Read the documents in this collection.
 
@@ -189,14 +194,14 @@ async def get(
             transaction
                 (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
                 An existing transaction that this query will run in.
-            retry (google.api_core.retry.Retry): Designation of what errors, if any,
-                should be retried.  Defaults to a system-specified policy.
-            timeout (float): The timeout for this request.  Defaults to a
-                system-specified value.
+            retry (Optional[google.api_core.retry.Retry]): Designation of what
+                errors, if any, should be retried.  Defaults to a
+                system-specified policy.
+            timeout (Otional[float]): The timeout for this request.  Defaults
+                to a system-specified value.
 
-        If a ``transaction`` is used and it already has write operations
-        added, this method cannot be used (i.e. read-after-write is not
-        allowed).
+        If a ``transaction`` is used and it already has write operations added,
+        this method cannot be used (i.e. read-after-write is not allowed).
 
         Returns:
             list: The documents in this collection that match the query.
@@ -205,15 +210,15 @@ async def get(
 
         return await query.get(transaction=transaction, **kwargs)
 
-    async def stream(
+    def stream(
         self,
-        transaction: Transaction = None,
-        retry: retries.AsyncRetry = gapic_v1.method.DEFAULT,
-        timeout: float = None,
-    ) -> AsyncIterator[async_document.DocumentSnapshot]:
+        transaction: Optional[transaction.Transaction] = None,
+        retry: Optional[retries.AsyncRetry] = gapic_v1.method.DEFAULT,
+        timeout: Optional[float] = None,
+    ) -> "AsyncStreamGenerator[DocumentSnapshot]":
         """Read the documents in this collection.
 
-        This sends a ``RunQuery`` RPC and then returns an iterator which
+        This sends a ``RunQuery`` RPC and then returns a generator which
         consumes each document returned in the stream of ``RunQueryResponse``
         messages.
 
@@ -232,16 +237,16 @@ async def stream(
             transaction (Optional[:class:`~google.cloud.firestore_v1.transaction.\
                 Transaction`]):
                 An existing transaction that the query will run in.
-            retry (google.api_core.retry.Retry): Designation of what errors, if any,
-                should be retried.  Defaults to a system-specified policy.
-            timeout (float): The timeout for this request.  Defaults to a
-                system-specified value.
+            retry (Optional[google.api_core.retry.Retry]): Designation of what
+                errors, if any, should be retried.  Defaults to a
+                system-specified policy.
+            timeout (Optional[float]): The timeout for this request. Defaults
+                to a system-specified value.
 
-        Yields:
-            :class:`~google.cloud.firestore_v1.document.DocumentSnapshot`:
-            The next document that fulfills the query.
+        Returns:
+            `AsyncStreamGenerator[DocumentSnapshot]`: A generator of the query
+            results.
         """
         query, kwargs = self._prep_get_or_stream(retry, timeout)
 
-        async for d in query.stream(transaction=transaction, **kwargs):
-            yield d  # pytype: disable=name-error
+        return query.stream(transaction=transaction, **kwargs)