Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/google/cloud/firestore_v1/async_aggregation.py: 53%

1# Copyright 2023 Google LLC All rights reserved. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""Classes for representing Async aggregation queries for the Google Cloud Firestore API. 

16 

17A :class:`~google.cloud.firestore_v1.async_aggregation.AsyncAggregationQuery` can be created directly from 

18a :class:`~google.cloud.firestore_v1.async_collection.AsyncCollection` and that can be 

19a more common way to create an aggregation query than direct usage of the constructor. 

20""" 

21from __future__ import annotations 

22 

23from google.api_core import gapic_v1 

24from google.api_core import retry as retries 

25 

26from typing import List, Union, AsyncGenerator 

27 

28 

29from google.cloud.firestore_v1.base_aggregation import ( 

30 AggregationResult, 

31 _query_response_to_result, 

32 BaseAggregationQuery, 

33) 

34 

35 

36class AsyncAggregationQuery(BaseAggregationQuery): 

37 """Represents an aggregation query to the Firestore API.""" 

38 

39 def __init__( 

40 self, 

41 nested_query, 

42 ) -> None: 

43 super(AsyncAggregationQuery, self).__init__(nested_query) 

44 

45 async def get( 

46 self, 

47 transaction=None, 

48 retry: Union[ 

49 retries.Retry, None, gapic_v1.method._MethodDefault 

50 ] = gapic_v1.method.DEFAULT, 

51 timeout: float | None = None, 

52 ) -> List[AggregationResult]: 

53 """Runs the aggregation query. 

54 

55 This sends a ``RunAggregationQuery`` RPC and returns a list of aggregation results in the stream of ``RunAggregationQueryResponse`` messages. 

56 

57 Args: 

58 transaction 

59 (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]): 

60 An existing transaction that this query will run in. 

61 If a ``transaction`` is used and it already has write operations 

62 added, this method cannot be used (i.e. read-after-write is not 

63 allowed). 

64 retry (google.api_core.retry.Retry): Designation of what errors, if any, 

65 should be retried. Defaults to a system-specified policy. 

66 timeout (float): The timeout for this request. Defaults to a 

67 system-specified value. 

68 

69 Returns: 

70 list: The aggregation query results 

71 

72 """ 

73 stream_result = self.stream( 

74 transaction=transaction, retry=retry, timeout=timeout 

75 ) 

76 result = [aggregation async for aggregation in stream_result] 

77 return result # type: ignore 

78 

79 async def stream( 

80 self, 

81 transaction=None, 

82 retry: Union[ 

83 retries.Retry, None, gapic_v1.method._MethodDefault 

84 ] = gapic_v1.method.DEFAULT, 

85 timeout: float | None = None, 

86 ) -> Union[AsyncGenerator[List[AggregationResult], None]]: 

87 """Runs the aggregation query. 

88 

89 This sends a ``RunAggregationQuery`` RPC and then returns an iterator which 

90 consumes each document returned in the stream of ``RunAggregationQueryResponse`` 

91 messages. 

92 

93 If a ``transaction`` is used and it already has write operations 

94 added, this method cannot be used (i.e. read-after-write is not 

95 allowed). 

96 

97 Args: 

98 transaction 

99 (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]): 

100 An existing transaction that this query will run in. 

101 retry (google.api_core.retry.Retry): Designation of what errors, if any, 

102 should be retried. Defaults to a system-specified policy. 

103 timeout (float): The timeout for this request. Defaults to a 

104 system-specified value. 

105 

106 Yields: 

107 :class:`~google.cloud.firestore_v1.base_aggregation.AggregationResult`: 

108 The result of aggregations of this query 

109 """ 

110 request, kwargs = self._prep_stream( 

111 transaction, 

112 retry, 

113 timeout, 

114 ) 

115 

116 response_iterator = await self._client._firestore_api.run_aggregation_query( 

117 request=request, 

118 metadata=self._client._rpc_metadata, 

119 **kwargs, 

120 ) 

121 

122 async for response in response_iterator: 

123 result = _query_response_to_result(response) 

124 yield result