This repository was archived by the owner on Mar 2, 2026. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 85
Expand file tree
/
Copy pathbase_aggregation.py
More file actions
376 lines (316 loc) · 14.5 KB
/
base_aggregation.py
File metadata and controls
376 lines (316 loc) · 14.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
# Copyright 2023 Google LLC All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Classes for representing aggregation queries for the Google Cloud Firestore API.
A :class:`~google.cloud.firestore_v1.aggregation.AggregationQuery` can be created directly from
a :class:`~google.cloud.firestore_v1.collection.Collection` and that can be
a more common way to create an aggregation query than direct usage of the constructor.
"""
from __future__ import annotations
import abc
import itertools
from abc import ABC
from typing import TYPE_CHECKING, Any, Coroutine, List, Optional, Tuple, Union, Iterable
from google.api_core import gapic_v1
from google.api_core import retry as retries
from google.cloud.firestore_v1 import _helpers
from google.cloud.firestore_v1.field_path import FieldPath
from google.cloud.firestore_v1.types import (
StructuredAggregationQuery,
)
from google.cloud.firestore_v1.pipeline_expressions import AggregateFunction
from google.cloud.firestore_v1.pipeline_expressions import Count
from google.cloud.firestore_v1.pipeline_expressions import AliasedExpression
from google.cloud.firestore_v1.pipeline_expressions import Field
# Types needed only for Type Hints
if TYPE_CHECKING: # pragma: NO COVER
from google.cloud.firestore_v1 import transaction
from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator
from google.cloud.firestore_v1.query_profile import ExplainOptions
from google.cloud.firestore_v1.query_results import QueryResultsList
from google.cloud.firestore_v1.stream_generator import (
StreamGenerator,
)
from google.cloud.firestore_v1.pipeline_source import PipelineSource
import datetime
class AggregationResult(object):
"""
A class representing result from Aggregation Query
:type alias: str
:param alias: The alias for the aggregation.
:type value: int
:param value: The resulting value from the aggregation.
:type read_time:
:param value: The resulting read_time
"""
def __init__(self, alias: str, value: float, read_time=None):
self.alias = alias
self.value = value
self.read_time = read_time
def __repr__(self):
return f"<Aggregation alias={self.alias}, value={self.value}, readtime={self.read_time}>"
def _to_dict(self):
return {self.alias: self.value}
class BaseAggregation(ABC):
def __init__(self, alias: str | None = None):
self.alias = alias
@abc.abstractmethod
def _to_protobuf(self):
"""Convert this instance to the protobuf representation"""
@abc.abstractmethod
def _to_pipeline_expr(
self, autoindexer: Iterable[int]
) -> AliasedExpression[AggregateFunction]:
"""
Convert this instance to a pipeline expression for use with pipeline.aggregate()
Args:
autoindexer: If an alias isn't supplied, one should be created with the format "field_n"
The autoindexer is an iterable that provides the `n` value to use for each expression
"""
def _pipeline_alias(self, autoindexer):
"""
Helper to build the alias for the pipeline expression
"""
if self.alias is not None:
return self.alias
else:
return f"field_{next(autoindexer)}"
class CountAggregation(BaseAggregation):
def __init__(self, alias: str | None = None):
super(CountAggregation, self).__init__(alias=alias)
def _to_protobuf(self):
"""Convert this instance to the protobuf representation"""
aggregation_pb = StructuredAggregationQuery.Aggregation()
if self.alias:
aggregation_pb.alias = self.alias
aggregation_pb.count = StructuredAggregationQuery.Aggregation.Count()
return aggregation_pb
def _to_pipeline_expr(self, autoindexer: Iterable[int]):
return Count().as_(self._pipeline_alias(autoindexer))
class SumAggregation(BaseAggregation):
def __init__(self, field_ref: str | FieldPath, alias: str | None = None):
# convert field path to string if needed
field_str = (
field_ref.to_api_repr() if isinstance(field_ref, FieldPath) else field_ref
)
self.field_ref: str = field_str
super(SumAggregation, self).__init__(alias=alias)
def _to_protobuf(self):
"""Convert this instance to the protobuf representation"""
aggregation_pb = StructuredAggregationQuery.Aggregation()
if self.alias:
aggregation_pb.alias = self.alias
aggregation_pb.sum = StructuredAggregationQuery.Aggregation.Sum()
aggregation_pb.sum.field.field_path = self.field_ref
return aggregation_pb
def _to_pipeline_expr(self, autoindexer: Iterable[int]):
return Field.of(self.field_ref).sum().as_(self._pipeline_alias(autoindexer))
class AvgAggregation(BaseAggregation):
def __init__(self, field_ref: str | FieldPath, alias: str | None = None):
# convert field path to string if needed
field_str = (
field_ref.to_api_repr() if isinstance(field_ref, FieldPath) else field_ref
)
self.field_ref: str = field_str
super(AvgAggregation, self).__init__(alias=alias)
def _to_protobuf(self):
"""Convert this instance to the protobuf representation"""
aggregation_pb = StructuredAggregationQuery.Aggregation()
if self.alias:
aggregation_pb.alias = self.alias
aggregation_pb.avg = StructuredAggregationQuery.Aggregation.Avg()
aggregation_pb.avg.field.field_path = self.field_ref
return aggregation_pb
def _to_pipeline_expr(self, autoindexer: Iterable[int]):
return Field.of(self.field_ref).average().as_(self._pipeline_alias(autoindexer))
def _query_response_to_result(
response_pb,
) -> List[AggregationResult]:
results = [
AggregationResult(
alias=key,
value=response_pb.result.aggregate_fields[key].integer_value
or response_pb.result.aggregate_fields[key].double_value,
read_time=response_pb.read_time,
)
for key in response_pb.result.aggregate_fields.pb.keys()
]
return results
class BaseAggregationQuery(ABC):
"""Represents an aggregation query to the Firestore API."""
def __init__(self, nested_query, alias: str | None = None) -> None:
self._nested_query = nested_query
self._alias = alias
self._collection_ref = nested_query._parent
self._aggregations: List[BaseAggregation] = []
@property
def _client(self):
return self._collection_ref._client
def count(self, alias: str | None = None):
"""
Adds a count over the nested query
"""
count_aggregation = CountAggregation(alias=alias)
self._aggregations.append(count_aggregation)
return self
def sum(self, field_ref: str | FieldPath, alias: str | None = None):
"""
Adds a sum over the nested query
"""
sum_aggregation = SumAggregation(field_ref, alias=alias)
self._aggregations.append(sum_aggregation)
return self
def avg(self, field_ref: str | FieldPath, alias: str | None = None):
"""
Adds an avg over the nested query
"""
avg_aggregation = AvgAggregation(field_ref, alias=alias)
self._aggregations.append(avg_aggregation)
return self
def add_aggregation(self, aggregation: BaseAggregation) -> None:
"""
Adds an aggregation operation to the nested query
:type aggregation: :class:`google.cloud.firestore_v1.aggregation.BaseAggregation`
:param aggregation: An aggregation operation, e.g. a CountAggregation
"""
self._aggregations.append(aggregation)
def add_aggregations(self, aggregations: List[BaseAggregation]) -> None:
"""
Adds a list of aggregations to the nested query
:type aggregations: list
:param aggregations: a list of aggregation operations
"""
self._aggregations.extend(aggregations)
def _to_protobuf(self) -> StructuredAggregationQuery:
pb = StructuredAggregationQuery()
pb.structured_query = self._nested_query._to_protobuf()
for aggregation in self._aggregations:
aggregation_pb = aggregation._to_protobuf()
pb.aggregations.append(aggregation_pb)
return pb
def _prep_stream(
self,
transaction=None,
retry: Union[retries.Retry, retries.AsyncRetry, None, object] = None,
timeout: float | None = None,
explain_options: Optional[ExplainOptions] = None,
read_time: Optional[datetime.datetime] = None,
) -> Tuple[dict, dict]:
parent_path, expected_prefix = self._collection_ref._parent_info()
request = {
"parent": parent_path,
"structured_aggregation_query": self._to_protobuf(),
"transaction": _helpers.get_transaction_id(transaction),
}
if explain_options:
request["explain_options"] = explain_options._to_dict()
if read_time is not None:
request["read_time"] = read_time
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
return request, kwargs
@abc.abstractmethod
def get(
self,
transaction=None,
retry: Union[
retries.Retry, retries.AsyncRetry, None, object
] = gapic_v1.method.DEFAULT,
timeout: float | None = None,
*,
explain_options: Optional[ExplainOptions] = None,
read_time: Optional[datetime.datetime] = None,
) -> (
QueryResultsList[AggregationResult]
| Coroutine[Any, Any, List[List[AggregationResult]]]
):
"""Runs the aggregation query.
This sends a ``RunAggregationQuery`` RPC and returns a list of
aggregation results in the stream of ``RunAggregationQueryResponse``
messages.
Args:
transaction
(Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
An existing transaction that this query will run in.
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).
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.
explain_options
(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.
read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.
Returns:
(QueryResultsList[List[AggregationResult]] | Coroutine[Any, Any, List[List[AggregationResult]]]):
The aggregation query results.
"""
@abc.abstractmethod
def stream(
self,
transaction: Optional[transaction.Transaction] = None,
retry: retries.Retry
| retries.AsyncRetry
| object
| None = gapic_v1.method.DEFAULT,
timeout: Optional[float] = None,
*,
explain_options: Optional[ExplainOptions] = None,
read_time: Optional[datetime.datetime] = None,
) -> (
StreamGenerator[List[AggregationResult]]
| AsyncStreamGenerator[List[AggregationResult]]
):
"""Runs the aggregation query.
This sends a``RunAggregationQuery`` RPC and returns a generator in the stream of ``RunAggregationQueryResponse`` messages.
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.
explain_options
(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.
read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.
Returns:
StreamGenerator[List[AggregationResult]] | AsyncStreamGenerator[List[AggregationResult]]:
A generator of the query results.
"""
def _build_pipeline(self, source: "PipelineSource"):
"""
Convert this query into a Pipeline
Queries containing a `cursor` or `limit_to_last` are not currently supported
Args:
source: the PipelineSource to build the pipeline off of
Raises:
- NotImplementedError: raised if the query contains a `cursor` or `limit_to_last`
Returns:
a Pipeline representing the query
"""
# use autoindexer to keep track of which field number to use for un-aliased fields
autoindexer = itertools.count(start=1)
exprs = [a._to_pipeline_expr(autoindexer) for a in self._aggregations]
return self._nested_query._build_pipeline(source).aggregate(*exprs)