diff --git a/google/cloud/firestore_v1/async_aggregation.py b/google/cloud/firestore_v1/async_aggregation.py index 479d82d61..dbaa27d95 100644 --- a/google/cloud/firestore_v1/async_aggregation.py +++ b/google/cloud/firestore_v1/async_aggregation.py @@ -132,7 +132,10 @@ async def _make_stream( Yields: List[AggregationResult] | query_profile_pb.ExplainMetrics: - The result of aggregations of this query. + The result of aggregations of this query. Query results will be + yielded as `List[AggregationResult]`. When the result contains + returned explain metrics, yield `query_profile_pb.ExplainMetrics` + individually. """ metrics: query_profile_pb.ExplainMetrics | None = None diff --git a/google/cloud/firestore_v1/async_query.py b/google/cloud/firestore_v1/async_query.py index b32e578d3..71b47964f 100644 --- a/google/cloud/firestore_v1/async_query.py +++ b/google/cloud/firestore_v1/async_query.py @@ -369,7 +369,9 @@ async def _make_stream( Yields: [:class:`~google.cloud.firestore_v1.base_document.DocumentSnapshot` \ | google.cloud.firestore_v1.types.query_profile.ExplainMetrtics]: - The next document that fulfills the query. + The next document that fulfills the query. Query results will be + yielded as `DocumentSnapshot`. When the result contains returned + explain metrics, yield `query_profile_pb.ExplainMetrics` individually. """ metrics: query_profile_pb.ExplainMetrics | None = None diff --git a/google/cloud/firestore_v1/async_transaction.py b/google/cloud/firestore_v1/async_transaction.py index bcc85b9b7..31a4e609c 100644 --- a/google/cloud/firestore_v1/async_transaction.py +++ b/google/cloud/firestore_v1/async_transaction.py @@ -173,7 +173,7 @@ async def get_all( async def get( self, - ref_or_query, + ref_or_query: AsyncDocumentReference | AsyncQuery, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, *, @@ -183,7 +183,8 @@ async def get( Retrieve a document or a query result from the database. Args: - ref_or_query The document references or query object to return. + ref_or_query (AsyncDocumentReference | AsyncQuery): + The document references or query object to return. 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 @@ -192,6 +193,8 @@ async def get( (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): Options to enable query profiling for this query. When set, explain_metrics will be available on the returned generator. + Can only be used when running a query, not a document reference. + Yields: DocumentSnapshot: The next document snapshot that fulfills the query, @@ -211,7 +214,7 @@ async def get( elif isinstance(ref_or_query, AsyncQuery): if explain_options is not None: kwargs["explain_options"] = explain_options - return ref_or_query.stream(transaction=self, **kwargs) + return await ref_or_query.stream(transaction=self, **kwargs) else: raise ValueError( 'Value for argument "ref_or_query" must be a AsyncDocumentReference or a AsyncQuery.' diff --git a/google/cloud/firestore_v1/async_vector_query.py b/google/cloud/firestore_v1/async_vector_query.py index 9b2032cbf..488ddc246 100644 --- a/google/cloud/firestore_v1/async_vector_query.py +++ b/google/cloud/firestore_v1/async_vector_query.py @@ -142,7 +142,9 @@ async def _make_stream( Yields: [:class:`~google.cloud.firestore_v1.base_document.DocumentSnapshot` \ | google.cloud.firestore_v1.types.query_profile.ExplainMetrtics]: - The next document that fulfills the query. + The next document that fulfills the query. Query results will be + yielded as `DocumentSnapshot`. When the result contains returned + explain metrics, yield `query_profile_pb.ExplainMetrics` individually. """ metrics: query_profile_pb.ExplainMetrics | None = None diff --git a/google/cloud/firestore_v1/transaction.py b/google/cloud/firestore_v1/transaction.py index 3f5edd363..0d6bf9d19 100644 --- a/google/cloud/firestore_v1/transaction.py +++ b/google/cloud/firestore_v1/transaction.py @@ -173,16 +173,17 @@ def get_all( def get( self, - ref_or_query, + ref_or_query: DocumentReference | Query, retry: retries.Retry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, *, explain_options: Optional[ExplainOptions] = None, ) -> StreamGenerator[DocumentSnapshot] | Generator[DocumentSnapshot, Any, None]: """Retrieve a document or a query result from the database. Args: - ref_or_query: The document references or query object to return. + ref_or_query (DocumentReference | Query): + The document references or query object to return. 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 @@ -191,7 +192,7 @@ def get( (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): Options to enable query profiling for this query. When set, explain_metrics will be available on the returned generator. - Can only be used when running a query. + Can only be used when running a query, not a document reference. Yields: .DocumentSnapshot: The next document snapshot that fulfills the diff --git a/tests/unit/v1/test_async_transaction.py b/tests/unit/v1/test_async_transaction.py index b685c2b7e..9dd495fb1 100644 --- a/tests/unit/v1/test_async_transaction.py +++ b/tests/unit/v1/test_async_transaction.py @@ -362,7 +362,7 @@ async def _get_w_query_helper(retry=None, timeout=None, explain_options=None): client = AsyncMock(spec=[]) transaction = _make_async_transaction(client) query = AsyncQuery(parent=AsyncMock(spec=[])) - query.stream = mock.Mock() + query.stream = AsyncMock() kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) result = await transaction.get(