1# Copyright 2017 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 queries for the Google Cloud Firestore API.
16
17A :class:`~google.cloud.firestore_v1.query.Query` can be created directly from
18a :class:`~google.cloud.firestore_v1.collection.Collection` and that can be
19a more common way to create a query than direct usage of the constructor.
20"""
21from __future__ import annotations
22
23import abc
24import copy
25import math
26import warnings
27
28from typing import (
29 TYPE_CHECKING,
30 Any,
31 Coroutine,
32 Dict,
33 Iterable,
34 List,
35 Optional,
36 Sequence,
37 Tuple,
38 Type,
39 Union,
40 TypeVar,
41)
42
43from google.api_core import retry as retries
44from google.protobuf import wrappers_pb2
45
46from google.cloud import firestore_v1
47from google.cloud.firestore_v1 import _helpers, document
48from google.cloud.firestore_v1 import field_path as field_path_module
49from google.cloud.firestore_v1 import transforms
50
51# Types needed only for Type Hints
52from google.cloud.firestore_v1.base_document import DocumentSnapshot
53from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
54from google.cloud.firestore_v1.order import Order
55from google.cloud.firestore_v1.types import (
56 Cursor,
57 RunQueryResponse,
58 StructuredQuery,
59 query,
60)
61from google.cloud.firestore_v1.vector import Vector
62
63if TYPE_CHECKING: # pragma: NO COVER
64 from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator
65 from google.cloud.firestore_v1.field_path import FieldPath
66 from google.cloud.firestore_v1.query_profile import ExplainOptions
67 from google.cloud.firestore_v1.query_results import QueryResultsList
68 from google.cloud.firestore_v1.stream_generator import StreamGenerator
69
70 import datetime
71
72
73_BAD_DIR_STRING: str
74_BAD_OP_NAN: str
75_BAD_OP_NULL: str
76_BAD_OP_STRING: str
77_COMPARISON_OPERATORS: Dict[str, Any]
78_EQ_OP: str
79_NEQ_OP: str
80_INVALID_CURSOR_TRANSFORM: str
81_INVALID_WHERE_TRANSFORM: str
82_MISMATCH_CURSOR_W_ORDER_BY: str
83_MISSING_ORDER_BY: str
84_NO_ORDERS_FOR_CURSOR: str
85_operator_enum: Any
86
87
88_EQ_OP = "=="
89_NEQ_OP = "!="
90_operator_enum = StructuredQuery.FieldFilter.Operator
91_COMPARISON_OPERATORS = {
92 "<": _operator_enum.LESS_THAN,
93 "<=": _operator_enum.LESS_THAN_OR_EQUAL,
94 _EQ_OP: _operator_enum.EQUAL,
95 _NEQ_OP: _operator_enum.NOT_EQUAL,
96 ">=": _operator_enum.GREATER_THAN_OR_EQUAL,
97 ">": _operator_enum.GREATER_THAN,
98 "array_contains": _operator_enum.ARRAY_CONTAINS,
99 "in": _operator_enum.IN,
100 "not-in": _operator_enum.NOT_IN,
101 "array_contains_any": _operator_enum.ARRAY_CONTAINS_ANY,
102}
103# set of operators that don't involve equlity comparisons
104# will be used in query normalization
105_INEQUALITY_OPERATORS = (
106 _operator_enum.LESS_THAN,
107 _operator_enum.LESS_THAN_OR_EQUAL,
108 _operator_enum.GREATER_THAN_OR_EQUAL,
109 _operator_enum.GREATER_THAN,
110 _operator_enum.NOT_EQUAL,
111 _operator_enum.NOT_IN,
112)
113_BAD_OP_STRING = "Operator string {!r} is invalid. Valid choices are: {}."
114_BAD_OP_NAN_NULL = 'Only equality ("==") or not-equal ("!=") filters can be used with None or NaN values'
115_INVALID_WHERE_TRANSFORM = "Transforms cannot be used as where values."
116_BAD_DIR_STRING = "Invalid direction {!r}. Must be one of {!r} or {!r}."
117_INVALID_CURSOR_TRANSFORM = "Transforms cannot be used as cursor values."
118_MISSING_ORDER_BY = (
119 'The "order by" field path {!r} is not present in the cursor data {!r}. '
120 "All fields sent to ``order_by()`` must be present in the fields "
121 "if passed to one of ``start_at()`` / ``start_after()`` / "
122 "``end_before()`` / ``end_at()`` to define a cursor."
123)
124
125_NO_ORDERS_FOR_CURSOR = (
126 "Attempting to create a cursor with no fields to order on. "
127 "When defining a cursor with one of ``start_at()`` / ``start_after()`` / "
128 "``end_before()`` / ``end_at()``, all fields in the cursor must "
129 "come from fields set in ``order_by()``."
130)
131_MISMATCH_CURSOR_W_ORDER_BY = "The cursor {!r} does not match the order fields {!r}."
132
133_not_passed = object()
134
135QueryType = TypeVar("QueryType", bound="BaseQuery")
136
137
138class BaseFilter(abc.ABC):
139 """Base class for Filters"""
140
141 @abc.abstractmethod
142 def _to_pb(self):
143 """Build the protobuf representation based on values in the filter"""
144
145
146def _validate_opation(op_string, value):
147 """
148 Given an input operator string (e.g, '!='), and a value (e.g. None),
149 ensure that the operator and value combination is valid, and return
150 an approproate new operator value. A new operator will be used if
151 the operaion is a comparison against Null or NaN
152
153 Args:
154 op_string (Optional[str]): the requested operator
155 value (Any): the value the operator is acting on
156 Returns:
157 str | StructuredQuery.UnaryFilter.Operator: operator to use in requests
158 Raises:
159 ValueError: if the operator and value combination is invalid
160 """
161 if value is None:
162 if op_string == _EQ_OP:
163 return StructuredQuery.UnaryFilter.Operator.IS_NULL
164 elif op_string == _NEQ_OP:
165 return StructuredQuery.UnaryFilter.Operator.IS_NOT_NULL
166 else:
167 raise ValueError(_BAD_OP_NAN_NULL)
168
169 elif _isnan(value):
170 if op_string == _EQ_OP:
171 return StructuredQuery.UnaryFilter.Operator.IS_NAN
172 elif op_string == _NEQ_OP:
173 return StructuredQuery.UnaryFilter.Operator.IS_NOT_NAN
174 else:
175 raise ValueError(_BAD_OP_NAN_NULL)
176 elif isinstance(value, (transforms.Sentinel, transforms._ValueList)):
177 raise ValueError(_INVALID_WHERE_TRANSFORM)
178 else:
179 return op_string
180
181
182class FieldFilter(BaseFilter):
183 """Class representation of a Field Filter."""
184
185 def __init__(self, field_path: str, op_string: str, value: Any | None = None):
186 self.field_path = field_path
187 self.value = value
188 self.op_string = _validate_opation(op_string, value)
189
190 def _to_pb(self):
191 """Returns the protobuf representation, either a StructuredQuery.UnaryFilter or a StructuredQuery.FieldFilter"""
192 if self.value is None or _isnan(self.value):
193 filter_pb = query.StructuredQuery.UnaryFilter(
194 field=query.StructuredQuery.FieldReference(field_path=self.field_path),
195 op=self.op_string,
196 )
197 else:
198 filter_pb = query.StructuredQuery.FieldFilter(
199 field=query.StructuredQuery.FieldReference(field_path=self.field_path),
200 op=_enum_from_op_string(self.op_string),
201 value=_helpers.encode_value(self.value),
202 )
203 return filter_pb
204
205
206class BaseCompositeFilter(BaseFilter):
207 """Base class for a Composite Filter. (either OR or AND)."""
208
209 def __init__(
210 self,
211 operator: int = StructuredQuery.CompositeFilter.Operator.OPERATOR_UNSPECIFIED,
212 filters: list[BaseFilter] | None = None,
213 ):
214 self.operator = operator
215 if filters is None:
216 self.filters = []
217 else:
218 self.filters = filters
219
220 def __repr__(self):
221 repr = f"op: {self.operator}\nFilters:"
222 for filter in self.filters:
223 repr += f"\n\t{filter}"
224 return repr
225
226 def _to_pb(self):
227 """Build the protobuf representation based on values in the Composite Filter."""
228 filter_pb = StructuredQuery.CompositeFilter(
229 op=self.operator,
230 )
231 for filter in self.filters:
232 if isinstance(filter, BaseCompositeFilter):
233 fb = query.StructuredQuery.Filter(composite_filter=filter._to_pb())
234 else:
235 fb = _filter_pb(filter._to_pb())
236 filter_pb.filters.append(fb)
237
238 return filter_pb
239
240
241class Or(BaseCompositeFilter):
242 """Class representation of an OR Filter."""
243
244 def __init__(self, filters: list[BaseFilter]):
245 super().__init__(
246 operator=StructuredQuery.CompositeFilter.Operator.OR, filters=filters
247 )
248
249
250class And(BaseCompositeFilter):
251 """Class representation of an AND Filter."""
252
253 def __init__(self, filters: list[BaseFilter]):
254 super().__init__(
255 operator=StructuredQuery.CompositeFilter.Operator.AND, filters=filters
256 )
257
258
259class BaseQuery(object):
260 """Represents a query to the Firestore API.
261
262 Instances of this class are considered immutable: all methods that
263 would modify an instance instead return a new instance.
264
265 Args:
266 parent (:class:`~google.cloud.firestore_v1.collection.CollectionReference`):
267 The collection that this query applies to.
268 projection (Optional[:class:`google.cloud.firestore_v1.\
269 query.StructuredQuery.Projection`]):
270 A projection of document fields to limit the query results to.
271 field_filters (Optional[Tuple[:class:`google.cloud.firestore_v1.\
272 query.StructuredQuery.FieldFilter`, ...]]):
273 The filters to be applied in the query.
274 orders (Optional[Tuple[:class:`google.cloud.firestore_v1.\
275 query.StructuredQuery.Order`, ...]]):
276 The "order by" entries to use in the query.
277 limit (Optional[int]):
278 The maximum number of documents the query is allowed to return.
279 limit_to_last (Optional[bool]):
280 Denotes whether a provided limit is applied to the end of the result set.
281 offset (Optional[int]):
282 The number of results to skip.
283 start_at (Optional[Tuple[dict, bool]]):
284 Two-tuple of :
285
286 * a mapping of fields. Any field that is present in this mapping
287 must also be present in ``orders``
288 * an ``after`` flag
289
290 The fields and the flag combine to form a cursor used as
291 a starting point in a query result set. If the ``after``
292 flag is :data:`True`, the results will start just after any
293 documents which have fields matching the cursor, otherwise
294 any matching documents will be included in the result set.
295 When the query is formed, the document values
296 will be used in the order given by ``orders``.
297 end_at (Optional[Tuple[dict, bool]]):
298 Two-tuple of:
299
300 * a mapping of fields. Any field that is present in this mapping
301 must also be present in ``orders``
302 * a ``before`` flag
303
304 The fields and the flag combine to form a cursor used as
305 an ending point in a query result set. If the ``before``
306 flag is :data:`True`, the results will end just before any
307 documents which have fields matching the cursor, otherwise
308 any matching documents will be included in the result set.
309 When the query is formed, the document values
310 will be used in the order given by ``orders``.
311 all_descendants (Optional[bool]):
312 When false, selects only collections that are immediate children
313 of the `parent` specified in the containing `RunQueryRequest`.
314 When true, selects all descendant collections.
315 recursive (Optional[bool]):
316 When true, returns all documents and all documents in any subcollections
317 below them. Defaults to false.
318 """
319
320 ASCENDING = "ASCENDING"
321 """str: Sort query results in ascending order on a field."""
322 DESCENDING = "DESCENDING"
323 """str: Sort query results in descending order on a field."""
324
325 def __init__(
326 self,
327 parent,
328 projection=None,
329 field_filters=(),
330 orders=(),
331 limit=None,
332 limit_to_last=False,
333 offset=None,
334 start_at=None,
335 end_at=None,
336 all_descendants=False,
337 recursive=False,
338 ) -> None:
339 self._parent = parent
340 self._projection = projection
341 self._field_filters = field_filters
342 self._orders = orders
343 self._limit = limit
344 self._limit_to_last = limit_to_last
345 self._offset = offset
346 self._start_at = start_at
347 self._end_at = end_at
348 self._all_descendants = all_descendants
349 self._recursive = recursive
350
351 def __eq__(self, other):
352 if not isinstance(other, self.__class__):
353 return NotImplemented
354 return (
355 self._parent == other._parent
356 and self._projection == other._projection
357 and self._field_filters == other._field_filters
358 and self._orders == other._orders
359 and self._limit == other._limit
360 and self._limit_to_last == other._limit_to_last
361 and self._offset == other._offset
362 and self._start_at == other._start_at
363 and self._end_at == other._end_at
364 and self._all_descendants == other._all_descendants
365 )
366
367 @property
368 def _client(self):
369 """The client of the parent collection.
370
371 Returns:
372 :class:`~google.cloud.firestore_v1.client.Client`:
373 The client that owns this query.
374 """
375 return self._parent._client
376
377 def select(self: QueryType, field_paths: Iterable[str]) -> QueryType:
378 """Project documents matching query to a limited set of fields.
379
380 See :meth:`~google.cloud.firestore_v1.client.Client.field_path` for
381 more information on **field paths**.
382
383 If the current query already has a projection set (i.e. has already
384 called :meth:`~google.cloud.firestore_v1.query.Query.select`), this
385 will overwrite it.
386
387 Args:
388 field_paths (Iterable[str, ...]): An iterable of field paths
389 (``.``-delimited list of field names) to use as a projection
390 of document fields in the query results.
391
392 Returns:
393 :class:`~google.cloud.firestore_v1.query.Query`:
394 A "projected" query. Acts as a copy of the current query,
395 modified with the newly added projection.
396 Raises:
397 ValueError: If any ``field_path`` is invalid.
398 """
399 field_paths = list(field_paths)
400 for field_path in field_paths:
401 field_path_module.split_field_path(field_path)
402
403 new_projection = query.StructuredQuery.Projection(
404 fields=[
405 query.StructuredQuery.FieldReference(field_path=field_path)
406 for field_path in field_paths
407 ]
408 )
409 return self._copy(projection=new_projection)
410
411 def _copy(
412 self: QueryType,
413 *,
414 projection: Optional[query.StructuredQuery.Projection] | object = _not_passed,
415 field_filters: Optional[Tuple[query.StructuredQuery.FieldFilter]]
416 | object = _not_passed,
417 orders: Optional[Tuple[query.StructuredQuery.Order]] | object = _not_passed,
418 limit: Optional[int] | object = _not_passed,
419 limit_to_last: Optional[bool] | object = _not_passed,
420 offset: Optional[int] | object = _not_passed,
421 start_at: Optional[Tuple[dict, bool]] | object = _not_passed,
422 end_at: Optional[Tuple[dict, bool]] | object = _not_passed,
423 all_descendants: Optional[bool] | object = _not_passed,
424 recursive: Optional[bool] | object = _not_passed,
425 ) -> QueryType:
426 return self.__class__(
427 self._parent,
428 projection=self._evaluate_param(projection, self._projection),
429 field_filters=self._evaluate_param(field_filters, self._field_filters),
430 orders=self._evaluate_param(orders, self._orders),
431 limit=self._evaluate_param(limit, self._limit),
432 limit_to_last=self._evaluate_param(limit_to_last, self._limit_to_last),
433 offset=self._evaluate_param(offset, self._offset),
434 start_at=self._evaluate_param(start_at, self._start_at),
435 end_at=self._evaluate_param(end_at, self._end_at),
436 all_descendants=self._evaluate_param(
437 all_descendants, self._all_descendants
438 ),
439 recursive=self._evaluate_param(recursive, self._recursive),
440 )
441
442 def _evaluate_param(self, value, fallback_value):
443 """Helper which allows `None` to be passed into `copy` and be set on the
444 copy instead of being misinterpreted as an unpassed parameter."""
445 return value if value is not _not_passed else fallback_value
446
447 def where(
448 self: QueryType,
449 field_path: Optional[str] = None,
450 op_string: Optional[str] = None,
451 value=None,
452 *,
453 filter=None,
454 ) -> QueryType:
455 """Filter the query on a field.
456
457 See :meth:`~google.cloud.firestore_v1.client.Client.field_path` for
458 more information on **field paths**.
459
460 Returns a new :class:`~google.cloud.firestore_v1.query.Query` that
461 filters on a specific field path, according to an operation (e.g.
462 ``==`` or "equals") and a particular value to be paired with that
463 operation.
464
465 Args:
466 field_path (Optional[str]): A field path (``.``-delimited list of
467 field names) for the field to filter on.
468 op_string (Optional[str]): A comparison operation in the form of a string.
469 Acceptable values are ``<``, ``<=``, ``==``, ``!=``, ``>=``, ``>``,
470 ``in``, ``not-in``, ``array_contains`` and ``array_contains_any``.
471 value (Any): The value to compare the field against in the filter.
472 If ``value`` is :data:`None` or a NaN, then ``==`` is the only
473 allowed operation.
474
475 Returns:
476 :class:`~google.cloud.firestore_v1.query.Query`:
477 A filtered query. Acts as a copy of the current query,
478 modified with the newly added filter.
479
480 Raises:
481 ValueError: If
482 * ``field_path`` is invalid.
483 * If ``value`` is a NaN or :data:`None` and ``op_string`` is not ``==``.
484 * FieldFilter was passed without using the filter keyword argument.
485 * `And` or `Or` was passed without using the filter keyword argument .
486 * Both the positional arguments and the keyword argument `filter` were passed.
487 """
488
489 if isinstance(field_path, FieldFilter):
490 raise ValueError(
491 "FieldFilter object must be passed using keyword argument 'filter'"
492 )
493 if isinstance(field_path, BaseCompositeFilter):
494 raise ValueError(
495 "'Or' and 'And' objects must be passed using keyword argument 'filter'"
496 )
497
498 field_path_module.split_field_path(field_path)
499 new_filters = self._field_filters
500
501 if field_path is not None and op_string is not None:
502 if filter is not None:
503 raise ValueError(
504 "Can't pass in both the positional arguments and 'filter' at the same time"
505 )
506 warnings.warn(
507 "Detected filter using positional arguments. Prefer using the 'filter' keyword argument instead.",
508 UserWarning,
509 stacklevel=2,
510 )
511 op = _validate_opation(op_string, value)
512 if isinstance(op, StructuredQuery.UnaryFilter.Operator):
513 filter_pb = query.StructuredQuery.UnaryFilter(
514 field=query.StructuredQuery.FieldReference(field_path=field_path),
515 op=op,
516 )
517 else:
518 filter_pb = query.StructuredQuery.FieldFilter(
519 field=query.StructuredQuery.FieldReference(field_path=field_path),
520 op=_enum_from_op_string(op_string),
521 value=_helpers.encode_value(value),
522 )
523
524 new_filters += (filter_pb,)
525 elif isinstance(filter, BaseFilter):
526 new_filters += (filter._to_pb(),)
527 else:
528 raise ValueError(
529 "Filter must be provided through positional arguments or the 'filter' keyword argument."
530 )
531 return self._copy(field_filters=new_filters)
532
533 @staticmethod
534 def _make_order(field_path, direction) -> StructuredQuery.Order:
535 """Helper for :meth:`order_by`."""
536 return query.StructuredQuery.Order(
537 field=query.StructuredQuery.FieldReference(field_path=field_path),
538 direction=_enum_from_direction(direction),
539 )
540
541 def order_by(
542 self: QueryType, field_path: str, direction: str = ASCENDING
543 ) -> QueryType:
544 """Modify the query to add an order clause on a specific field.
545
546 See :meth:`~google.cloud.firestore_v1.client.Client.field_path` for
547 more information on **field paths**.
548
549 Successive :meth:`~google.cloud.firestore_v1.query.Query.order_by`
550 calls will further refine the ordering of results returned by the query
551 (i.e. the new "order by" fields will be added to existing ones).
552
553 Args:
554 field_path (str): A field path (``.``-delimited list of
555 field names) on which to order the query results.
556 direction (Optional[str]): The direction to order by. Must be one
557 of :attr:`ASCENDING` or :attr:`DESCENDING`, defaults to
558 :attr:`ASCENDING`.
559
560 Returns:
561 :class:`~google.cloud.firestore_v1.query.Query`:
562 An ordered query. Acts as a copy of the current query, modified
563 with the newly added "order by" constraint.
564
565 Raises:
566 ValueError: If ``field_path`` is invalid.
567 ValueError: If ``direction`` is not one of :attr:`ASCENDING` or
568 :attr:`DESCENDING`.
569 """
570 field_path_module.split_field_path(field_path) # raises
571
572 order_pb = self._make_order(field_path, direction)
573
574 new_orders = self._orders + (order_pb,)
575 return self._copy(orders=new_orders)
576
577 def limit(self: QueryType, count: int) -> QueryType:
578 """Limit a query to return at most `count` matching results.
579
580 If the current query already has a `limit` set, this will override it.
581
582 .. note::
583 `limit` and `limit_to_last` are mutually exclusive.
584 Setting `limit` will drop previously set `limit_to_last`.
585
586 Args:
587 count (int): Maximum number of documents to return that match
588 the query.
589 Returns:
590 :class:`~google.cloud.firestore_v1.query.Query`:
591 A limited query. Acts as a copy of the current query, modified
592 with the newly added "limit" filter.
593 """
594 return self._copy(limit=count, limit_to_last=False)
595
596 def limit_to_last(self: QueryType, count: int) -> QueryType:
597 """Limit a query to return the last `count` matching results.
598 If the current query already has a `limit_to_last`
599 set, this will override it.
600
601 .. note::
602 `limit` and `limit_to_last` are mutually exclusive.
603 Setting `limit_to_last` will drop previously set `limit`.
604
605 Args:
606 count (int): Maximum number of documents to return that match
607 the query.
608 Returns:
609 :class:`~google.cloud.firestore_v1.query.Query`:
610 A limited query. Acts as a copy of the current query, modified
611 with the newly added "limit" filter.
612 """
613 return self._copy(limit=count, limit_to_last=True)
614
615 def _resolve_chunk_size(self, num_loaded: int, chunk_size: int) -> int:
616 """Utility function for chunkify."""
617 if self._limit is not None and (num_loaded + chunk_size) > self._limit:
618 return max(self._limit - num_loaded, 0)
619 return chunk_size
620
621 def offset(self: QueryType, num_to_skip: int) -> QueryType:
622 """Skip to an offset in a query.
623
624 If the current query already has specified an offset, this will
625 overwrite it.
626
627 Args:
628 num_to_skip (int): The number of results to skip at the beginning
629 of query results. (Must be non-negative.)
630
631 Returns:
632 :class:`~google.cloud.firestore_v1.query.Query`:
633 An offset query. Acts as a copy of the current query, modified
634 with the newly added "offset" field.
635 """
636 return self._copy(offset=num_to_skip)
637
638 def _check_snapshot(self, document_snapshot) -> None:
639 """Validate local snapshots for non-collection-group queries.
640
641 Raises:
642 ValueError: for non-collection-group queries, if the snapshot
643 is from a different collection.
644 """
645 if self._all_descendants:
646 return
647
648 if document_snapshot.reference._path[:-1] != self._parent._path:
649 raise ValueError("Cannot use snapshot from another collection as a cursor.")
650
651 def _cursor_helper(
652 self: QueryType,
653 document_fields_or_snapshot: Union[DocumentSnapshot, dict, list, tuple, None],
654 before: bool,
655 start: bool,
656 ) -> QueryType:
657 """Set values to be used for a ``start_at`` or ``end_at`` cursor.
658
659 The values will later be used in a query protobuf.
660
661 When the query is sent to the server, the ``document_fields_or_snapshot`` will
662 be used in the order given by fields set by
663 :meth:`~google.cloud.firestore_v1.query.Query.order_by`.
664
665 Args:
666 document_fields_or_snapshot
667 (Union[:class:`~google.cloud.firestore_v1.document.DocumentSnapshot`, dict, list, tuple]):
668 a document snapshot or a dictionary/list/tuple of fields
669 representing a query results cursor. A cursor is a collection
670 of values that represent a position in a query result set.
671 before (bool): Flag indicating if the document in
672 ``document_fields_or_snapshot`` should (:data:`False`) or
673 shouldn't (:data:`True`) be included in the result set.
674 start (Optional[bool]): determines if the cursor is a ``start_at``
675 cursor (:data:`True`) or an ``end_at`` cursor (:data:`False`).
676
677 Returns:
678 :class:`~google.cloud.firestore_v1.query.Query`:
679 A query with cursor. Acts as a copy of the current query, modified
680 with the newly added "start at" cursor.
681 """
682 if isinstance(document_fields_or_snapshot, tuple):
683 document_fields_or_snapshot = list(document_fields_or_snapshot)
684 elif isinstance(document_fields_or_snapshot, document.DocumentSnapshot):
685 self._check_snapshot(document_fields_or_snapshot)
686 else:
687 # NOTE: We copy so that the caller can't modify after calling.
688 document_fields_or_snapshot = copy.deepcopy(document_fields_or_snapshot)
689
690 cursor_pair = document_fields_or_snapshot, before
691 query_kwargs = {
692 "projection": self._projection,
693 "field_filters": self._field_filters,
694 "orders": self._orders,
695 "limit": self._limit,
696 "offset": self._offset,
697 "all_descendants": self._all_descendants,
698 }
699 if start:
700 query_kwargs["start_at"] = cursor_pair
701 query_kwargs["end_at"] = self._end_at
702 else:
703 query_kwargs["start_at"] = self._start_at
704 query_kwargs["end_at"] = cursor_pair
705
706 return self._copy(**query_kwargs)
707
708 def start_at(
709 self: QueryType,
710 document_fields_or_snapshot: Union[DocumentSnapshot, dict, list, tuple, None],
711 ) -> QueryType:
712 """Start query results at a particular document value.
713
714 The result set will **include** the document specified by
715 ``document_fields_or_snapshot``.
716
717 If the current query already has specified a start cursor -- either
718 via this method or
719 :meth:`~google.cloud.firestore_v1.query.Query.start_after` -- this
720 will overwrite it.
721
722 When the query is sent to the server, the ``document_fields`` will
723 be used in the order given by fields set by
724 :meth:`~google.cloud.firestore_v1.query.Query.order_by`.
725
726 Args:
727 document_fields_or_snapshot
728 (Union[:class:`~google.cloud.firestore_v1.document.DocumentSnapshot`, dict, list, tuple]):
729 a document snapshot or a dictionary/list/tuple of fields
730 representing a query results cursor. A cursor is a collection
731 of values that represent a position in a query result set.
732
733 Returns:
734 :class:`~google.cloud.firestore_v1.query.Query`:
735 A query with cursor. Acts as
736 a copy of the current query, modified with the newly added
737 "start at" cursor.
738 """
739 return self._cursor_helper(document_fields_or_snapshot, before=True, start=True)
740
741 def start_after(
742 self: QueryType,
743 document_fields_or_snapshot: Union[DocumentSnapshot, dict, list, tuple, None],
744 ) -> QueryType:
745 """Start query results after a particular document value.
746
747 The result set will **exclude** the document specified by
748 ``document_fields_or_snapshot``.
749
750 If the current query already has specified a start cursor -- either
751 via this method or
752 :meth:`~google.cloud.firestore_v1.query.Query.start_at` -- this will
753 overwrite it.
754
755 When the query is sent to the server, the ``document_fields_or_snapshot`` will
756 be used in the order given by fields set by
757 :meth:`~google.cloud.firestore_v1.query.Query.order_by`.
758
759 Args:
760 document_fields_or_snapshot
761 (Union[:class:`~google.cloud.firestore_v1.document.DocumentSnapshot`, dict, list, tuple]):
762 a document snapshot or a dictionary/list/tuple of fields
763 representing a query results cursor. A cursor is a collection
764 of values that represent a position in a query result set.
765
766 Returns:
767 :class:`~google.cloud.firestore_v1.query.Query`:
768 A query with cursor. Acts as a copy of the current query, modified
769 with the newly added "start after" cursor.
770 """
771 return self._cursor_helper(
772 document_fields_or_snapshot, before=False, start=True
773 )
774
775 def end_before(
776 self: QueryType,
777 document_fields_or_snapshot: Union[DocumentSnapshot, dict, list, tuple, None],
778 ) -> QueryType:
779 """End query results before a particular document value.
780
781 The result set will **exclude** the document specified by
782 ``document_fields_or_snapshot``.
783
784 If the current query already has specified an end cursor -- either
785 via this method or
786 :meth:`~google.cloud.firestore_v1.query.Query.end_at` -- this will
787 overwrite it.
788
789 When the query is sent to the server, the ``document_fields_or_snapshot`` will
790 be used in the order given by fields set by
791 :meth:`~google.cloud.firestore_v1.query.Query.order_by`.
792
793 Args:
794 document_fields_or_snapshot
795 (Union[:class:`~google.cloud.firestore_v1.document.DocumentSnapshot`, dict, list, tuple]):
796 a document snapshot or a dictionary/list/tuple of fields
797 representing a query results cursor. A cursor is a collection
798 of values that represent a position in a query result set.
799
800 Returns:
801 :class:`~google.cloud.firestore_v1.query.Query`:
802 A query with cursor. Acts as a copy of the current query, modified
803 with the newly added "end before" cursor.
804 """
805 return self._cursor_helper(
806 document_fields_or_snapshot, before=True, start=False
807 )
808
809 def end_at(
810 self: QueryType,
811 document_fields_or_snapshot: Union[DocumentSnapshot, dict, list, tuple, None],
812 ) -> QueryType:
813 """End query results at a particular document value.
814
815 The result set will **include** the document specified by
816 ``document_fields_or_snapshot``.
817
818 If the current query already has specified an end cursor -- either
819 via this method or
820 :meth:`~google.cloud.firestore_v1.query.Query.end_before` -- this will
821 overwrite it.
822
823 When the query is sent to the server, the ``document_fields_or_snapshot`` will
824 be used in the order given by fields set by
825 :meth:`~google.cloud.firestore_v1.query.Query.order_by`.
826
827 Args:
828 document_fields_or_snapshot
829 (Union[:class:`~google.cloud.firestore_v1.document.DocumentSnapshot`, dict, list, tuple]):
830 a document snapshot or a dictionary/list/tuple of fields
831 representing a query results cursor. A cursor is a collection
832 of values that represent a position in a query result set.
833
834 Returns:
835 :class:`~google.cloud.firestore_v1.query.Query`:
836 A query with cursor. Acts as a copy of the current query, modified
837 with the newly added "end at" cursor.
838 """
839 return self._cursor_helper(
840 document_fields_or_snapshot, before=False, start=False
841 )
842
843 def _filters_pb(self) -> Optional[StructuredQuery.Filter]:
844 """Convert all the filters into a single generic Filter protobuf.
845
846 This may be a lone field filter or unary filter, may be a composite
847 filter or may be :data:`None`.
848
849 Returns:
850 :class:`google.cloud.firestore_v1.types.StructuredQuery.Filter`:
851 A "generic" filter representing the current query's filters.
852 """
853 num_filters = len(self._field_filters)
854 if num_filters == 0:
855 return None
856 elif num_filters == 1:
857 filter = self._field_filters[0]
858 if isinstance(filter, query.StructuredQuery.CompositeFilter):
859 return query.StructuredQuery.Filter(composite_filter=filter)
860 else:
861 return _filter_pb(filter)
862 else:
863 composite_filter = query.StructuredQuery.CompositeFilter(
864 op=StructuredQuery.CompositeFilter.Operator.AND,
865 )
866 for filter_ in self._field_filters:
867 if isinstance(filter_, query.StructuredQuery.CompositeFilter):
868 composite_filter.filters.append(
869 query.StructuredQuery.Filter(composite_filter=filter_)
870 )
871 else:
872 composite_filter.filters.append(_filter_pb(filter_))
873
874 return query.StructuredQuery.Filter(composite_filter=composite_filter)
875
876 @staticmethod
877 def _normalize_projection(projection) -> StructuredQuery.Projection:
878 """Helper: convert field paths to message."""
879 if projection is not None:
880 fields = list(projection.fields)
881
882 if not fields:
883 field_ref = query.StructuredQuery.FieldReference(field_path="__name__")
884 return query.StructuredQuery.Projection(fields=[field_ref])
885
886 return projection
887
888 def _normalize_orders(self) -> list:
889 """Helper: adjust orders based on cursors, where clauses."""
890 orders = list(self._orders)
891 _has_snapshot_cursor = False
892
893 if self._start_at:
894 if isinstance(self._start_at[0], document.DocumentSnapshot):
895 _has_snapshot_cursor = True
896
897 if self._end_at:
898 if isinstance(self._end_at[0], document.DocumentSnapshot):
899 _has_snapshot_cursor = True
900 if _has_snapshot_cursor:
901 # added orders should use direction of last order
902 last_direction = orders[-1].direction if orders else BaseQuery.ASCENDING
903 order_keys = [order.field.field_path for order in orders]
904 for filter_ in self._field_filters:
905 # FieldFilter.Operator should not compare equal to
906 # UnaryFilter.Operator, but it does
907 if isinstance(filter_.op, StructuredQuery.FieldFilter.Operator):
908 field = filter_.field.field_path
909 # skip equality filters and filters on fields already ordered
910 if filter_.op in _INEQUALITY_OPERATORS and field not in order_keys:
911 orders.append(self._make_order(field, last_direction))
912 # add __name__ if not already in orders
913 if "__name__" not in [order.field.field_path for order in orders]:
914 orders.append(self._make_order("__name__", last_direction))
915
916 return orders
917
918 def _normalize_cursor(self, cursor, orders) -> Tuple[List, bool] | None:
919 """Helper: convert cursor to a list of values based on orders."""
920 if cursor is None:
921 return None
922
923 if not orders:
924 raise ValueError(_NO_ORDERS_FOR_CURSOR)
925
926 document_fields, before = cursor
927
928 order_keys = [order.field.field_path for order in orders]
929
930 if isinstance(document_fields, document.DocumentSnapshot):
931 snapshot = document_fields
932 document_fields = snapshot.to_dict()
933 document_fields["__name__"] = snapshot.reference
934
935 if isinstance(document_fields, dict):
936 # Transform to list using orders
937 values = []
938 data = document_fields
939
940 # It isn't required that all order by have a cursor.
941 # However, we need to be sure they are specified in order without gaps
942 for order_key in order_keys[: len(data)]:
943 try:
944 if order_key in data:
945 values.append(data[order_key])
946 else:
947 values.append(
948 field_path_module.get_nested_value(order_key, data)
949 )
950 except KeyError:
951 msg = _MISSING_ORDER_BY.format(order_key, data)
952 raise ValueError(msg)
953
954 document_fields = values
955
956 if document_fields and len(document_fields) > len(orders):
957 msg = _MISMATCH_CURSOR_W_ORDER_BY.format(document_fields, order_keys)
958 raise ValueError(msg)
959
960 _transform_bases = (transforms.Sentinel, transforms._ValueList)
961
962 for index, key_field in enumerate(zip(order_keys, document_fields)):
963 key, field = key_field
964
965 if isinstance(field, _transform_bases):
966 msg = _INVALID_CURSOR_TRANSFORM
967 raise ValueError(msg)
968
969 if key == "__name__" and isinstance(field, str):
970 document_fields[index] = self._parent.document(field)
971
972 return document_fields, before
973
974 def _to_protobuf(self) -> StructuredQuery:
975 """Convert the current query into the equivalent protobuf.
976
977 Returns:
978 :class:`google.cloud.firestore_v1.types.StructuredQuery`:
979 The query protobuf.
980 """
981 projection = self._normalize_projection(self._projection)
982 orders = self._normalize_orders()
983 start_at = self._normalize_cursor(self._start_at, orders)
984 end_at = self._normalize_cursor(self._end_at, orders)
985
986 query_kwargs = {
987 "select": projection,
988 "from_": [
989 query.StructuredQuery.CollectionSelector(
990 collection_id=self._parent.id, all_descendants=self._all_descendants
991 )
992 ],
993 "where": self._filters_pb(),
994 "order_by": orders,
995 "start_at": _cursor_pb(start_at),
996 "end_at": _cursor_pb(end_at),
997 }
998 if self._offset is not None:
999 query_kwargs["offset"] = self._offset
1000 if self._limit is not None:
1001 query_kwargs["limit"] = wrappers_pb2.Int32Value(value=self._limit)
1002 return query.StructuredQuery(**query_kwargs)
1003
1004 def find_nearest(
1005 self,
1006 vector_field: str,
1007 query_vector: Union[Vector, Sequence[float]],
1008 limit: int,
1009 distance_measure: DistanceMeasure,
1010 *,
1011 distance_result_field: Optional[str] = None,
1012 distance_threshold: Optional[float] = None,
1013 ):
1014 raise NotImplementedError
1015
1016 def count(
1017 self, alias: str | None = None
1018 ) -> Type["firestore_v1.base_aggregation.BaseAggregationQuery"]:
1019 raise NotImplementedError
1020
1021 def sum(
1022 self, field_ref: str | FieldPath, alias: str | None = None
1023 ) -> Type["firestore_v1.base_aggregation.BaseAggregationQuery"]:
1024 raise NotImplementedError
1025
1026 def avg(
1027 self, field_ref: str | FieldPath, alias: str | None = None
1028 ) -> Type["firestore_v1.base_aggregation.BaseAggregationQuery"]:
1029 raise NotImplementedError
1030
1031 def get(
1032 self,
1033 transaction=None,
1034 retry: Optional[retries.Retry] = None,
1035 timeout: Optional[float] = None,
1036 *,
1037 explain_options: Optional[ExplainOptions] = None,
1038 read_time: Optional[datetime.datetime] = None,
1039 ) -> (
1040 QueryResultsList[DocumentSnapshot]
1041 | Coroutine[Any, Any, QueryResultsList[DocumentSnapshot]]
1042 ):
1043 raise NotImplementedError
1044
1045 def _prep_stream(
1046 self,
1047 transaction=None,
1048 retry: retries.Retry | retries.AsyncRetry | object | None = None,
1049 timeout: Optional[float] = None,
1050 explain_options: Optional[ExplainOptions] = None,
1051 read_time: Optional[datetime.datetime] = None,
1052 ) -> Tuple[dict, str, dict]:
1053 """Shared setup for async / sync :meth:`stream`"""
1054 if self._limit_to_last:
1055 raise ValueError(
1056 "Query results for queries that include limit_to_last() "
1057 "constraints cannot be streamed. Use Query.get() instead."
1058 )
1059
1060 parent_path, expected_prefix = self._parent._parent_info()
1061 request = {
1062 "parent": parent_path,
1063 "structured_query": self._to_protobuf(),
1064 "transaction": _helpers.get_transaction_id(transaction),
1065 }
1066 if explain_options is not None:
1067 request["explain_options"] = explain_options._to_dict()
1068 if read_time is not None:
1069 request["read_time"] = read_time
1070 kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
1071
1072 return request, expected_prefix, kwargs
1073
1074 def stream(
1075 self,
1076 transaction=None,
1077 retry: Optional[retries.Retry] = None,
1078 timeout: Optional[float] = None,
1079 *,
1080 explain_options: Optional[ExplainOptions] = None,
1081 read_time: Optional[datetime.datetime] = None,
1082 ) -> (
1083 StreamGenerator[document.DocumentSnapshot]
1084 | AsyncStreamGenerator[DocumentSnapshot]
1085 ):
1086 raise NotImplementedError
1087
1088 def on_snapshot(self, callback):
1089 raise NotImplementedError
1090
1091 def recursive(self: QueryType) -> QueryType:
1092 """Returns a copy of this query whose iterator will yield all matching
1093 documents as well as each of their descendent subcollections and documents.
1094
1095 This differs from the `all_descendents` flag, which only returns descendents
1096 whose subcollection names match the parent collection's name. To return
1097 all descendents, regardless of their subcollection name, use this.
1098 """
1099 copied = self._copy(recursive=True, all_descendants=True)
1100 if copied._parent and copied._parent.id:
1101 original_collection_id = "/".join(copied._parent._path)
1102
1103 # Reset the parent to nothing so we can recurse through the entire
1104 # database. This is required to have
1105 # `CollectionSelector.collection_id` not override
1106 # `CollectionSelector.all_descendants`, which happens if both are
1107 # set.
1108 copied._parent = copied._get_collection_reference_class()("")
1109 copied._parent._client = self._parent._client
1110
1111 # But wait! We don't want to load the entire database; only the
1112 # collection the user originally specified. To accomplish that, we
1113 # add the following arcane filters.
1114
1115 REFERENCE_NAME_MIN_ID = "__id-9223372036854775808__"
1116 start_at = f"{original_collection_id}/{REFERENCE_NAME_MIN_ID}"
1117
1118 # The backend interprets this null character is flipping the filter
1119 # to mean the end of the range instead of the beginning.
1120 nullChar = "\0"
1121 end_at = f"{original_collection_id}{nullChar}/{REFERENCE_NAME_MIN_ID}"
1122
1123 copied = (
1124 copied.order_by(field_path_module.FieldPath.document_id())
1125 .start_at({field_path_module.FieldPath.document_id(): start_at})
1126 .end_at({field_path_module.FieldPath.document_id(): end_at})
1127 )
1128
1129 return copied
1130
1131 def _comparator(self, doc1, doc2) -> int:
1132 _orders = self._orders
1133
1134 # Add implicit sorting by name, using the last specified direction.
1135 if len(_orders) == 0:
1136 lastDirection = BaseQuery.ASCENDING
1137 else:
1138 if _orders[-1].direction == 1:
1139 lastDirection = BaseQuery.ASCENDING
1140 else:
1141 lastDirection = BaseQuery.DESCENDING
1142
1143 orderBys = list(_orders)
1144
1145 order_pb = query.StructuredQuery.Order(
1146 field=query.StructuredQuery.FieldReference(field_path="id"),
1147 direction=_enum_from_direction(lastDirection),
1148 )
1149 orderBys.append(order_pb)
1150
1151 for orderBy in orderBys:
1152 if orderBy.field.field_path == "id":
1153 # If ordering by document id, compare resource paths.
1154 comp = Order()._compare_to(doc1.reference._path, doc2.reference._path)
1155 else:
1156 if (
1157 orderBy.field.field_path not in doc1._data
1158 or orderBy.field.field_path not in doc2._data
1159 ):
1160 raise ValueError(
1161 "Can only compare fields that exist in the "
1162 "DocumentSnapshot. Please include the fields you are "
1163 "ordering on in your select() call."
1164 )
1165 v1 = doc1._data[orderBy.field.field_path]
1166 v2 = doc2._data[orderBy.field.field_path]
1167 encoded_v1 = _helpers.encode_value(v1)
1168 encoded_v2 = _helpers.encode_value(v2)
1169 comp = Order().compare(encoded_v1, encoded_v2)
1170
1171 if comp != 0:
1172 # 1 == Ascending, -1 == Descending
1173 return orderBy.direction * comp
1174
1175 return 0
1176
1177 @staticmethod
1178 def _get_collection_reference_class():
1179 raise NotImplementedError
1180
1181
1182def _enum_from_op_string(op_string: str) -> int:
1183 """Convert a string representation of a binary operator to an enum.
1184
1185 These enums come from the protobuf message definition
1186 ``StructuredQuery.FieldFilter.Operator``.
1187
1188 Args:
1189 op_string (str): A comparison operation in the form of a string.
1190 Acceptable values are ``<``, ``<=``, ``==``, ``!=``, ``>=``
1191 and ``>``.
1192
1193 Returns:
1194 int: The enum corresponding to ``op_string``.
1195
1196 Raises:
1197 ValueError: If ``op_string`` is not a valid operator.
1198 """
1199 try:
1200 return _COMPARISON_OPERATORS[op_string]
1201 except KeyError:
1202 choices = ", ".join(sorted(_COMPARISON_OPERATORS.keys()))
1203 msg = _BAD_OP_STRING.format(op_string, choices)
1204 raise ValueError(msg)
1205
1206
1207def _isnan(value) -> bool:
1208 """Check if a value is NaN.
1209
1210 This differs from ``math.isnan`` in that **any** input type is
1211 allowed.
1212
1213 Args:
1214 value (Any): A value to check for NaN-ness.
1215
1216 Returns:
1217 bool: Indicates if the value is the NaN float.
1218 """
1219 if isinstance(value, float):
1220 return math.isnan(value)
1221 else:
1222 return False
1223
1224
1225def _enum_from_direction(direction: str) -> int:
1226 """Convert a string representation of a direction to an enum.
1227
1228 Args:
1229 direction (str): A direction to order by. Must be one of
1230 :attr:`~google.cloud.firestore.BaseQuery.ASCENDING` or
1231 :attr:`~google.cloud.firestore.BaseQuery.DESCENDING`.
1232
1233 Returns:
1234 int: The enum corresponding to ``direction``.
1235
1236 Raises:
1237 ValueError: If ``direction`` is not a valid direction.
1238 """
1239 if isinstance(direction, int):
1240 return direction
1241
1242 if direction == BaseQuery.ASCENDING:
1243 return StructuredQuery.Direction.ASCENDING
1244 elif direction == BaseQuery.DESCENDING:
1245 return StructuredQuery.Direction.DESCENDING
1246 else:
1247 msg = _BAD_DIR_STRING.format(
1248 direction, BaseQuery.ASCENDING, BaseQuery.DESCENDING
1249 )
1250 raise ValueError(msg)
1251
1252
1253def _filter_pb(field_or_unary) -> StructuredQuery.Filter:
1254 """Convert a specific protobuf filter to the generic filter type.
1255
1256 Args:
1257 field_or_unary (Union[google.cloud.firestore_v1.\
1258 query.StructuredQuery.FieldFilter, google.cloud.\
1259 firestore_v1.query.StructuredQuery.FieldFilter]): A
1260 field or unary filter to convert to a generic filter.
1261
1262 Returns:
1263 google.cloud.firestore_v1.types.\
1264 StructuredQuery.Filter: A "generic" filter.
1265
1266 Raises:
1267 ValueError: If ``field_or_unary`` is not a field or unary filter.
1268 """
1269 if isinstance(field_or_unary, query.StructuredQuery.FieldFilter):
1270 return query.StructuredQuery.Filter(field_filter=field_or_unary)
1271 elif isinstance(field_or_unary, query.StructuredQuery.UnaryFilter):
1272 return query.StructuredQuery.Filter(unary_filter=field_or_unary)
1273 else:
1274 raise ValueError("Unexpected filter type", type(field_or_unary), field_or_unary)
1275
1276
1277def _cursor_pb(cursor_pair: Optional[Tuple[list, bool]]) -> Optional[Cursor]:
1278 """Convert a cursor pair to a protobuf.
1279
1280 If ``cursor_pair`` is :data:`None`, just returns :data:`None`.
1281
1282 Args:
1283 cursor_pair (Optional[Tuple[list, bool]]): Two-tuple of
1284
1285 * a list of field values.
1286 * a ``before`` flag
1287
1288 Returns:
1289 Optional[google.cloud.firestore_v1.types.Cursor]: A
1290 protobuf cursor corresponding to the values.
1291 """
1292 if cursor_pair is not None:
1293 data, before = cursor_pair
1294 value_pbs = [_helpers.encode_value(value) for value in data]
1295 return query.Cursor(values=value_pbs, before=before)
1296 else:
1297 return None
1298
1299
1300def _query_response_to_snapshot(
1301 response_pb: RunQueryResponse, collection, expected_prefix: str
1302) -> Optional[document.DocumentSnapshot]:
1303 """Parse a query response protobuf to a document snapshot.
1304
1305 Args:
1306 response_pb (google.cloud.firestore_v1.\
1307 firestore.RunQueryResponse): A
1308 collection (:class:`~google.cloud.firestore_v1.collection.CollectionReference`):
1309 A reference to the collection that initiated the query.
1310 expected_prefix (str): The expected prefix for fully-qualified
1311 document names returned in the query results. This can be computed
1312 directly from ``collection`` via :meth:`_parent_info`.
1313
1314 Returns:
1315 Optional[:class:`~google.cloud.firestore.document.DocumentSnapshot`]:
1316 A snapshot of the data returned in the query. If
1317 ``response_pb.document`` is not set, the snapshot will be :data:`None`.
1318 """
1319 if not response_pb._pb.HasField("document"):
1320 return None
1321
1322 document_id = _helpers.get_doc_id(response_pb.document, expected_prefix)
1323 reference = collection.document(document_id)
1324 data = _helpers.decode_dict(response_pb.document.fields, collection._client)
1325 snapshot = document.DocumentSnapshot(
1326 reference,
1327 data,
1328 exists=True,
1329 read_time=response_pb.read_time,
1330 create_time=response_pb.document.create_time,
1331 update_time=response_pb.document.update_time,
1332 )
1333 return snapshot
1334
1335
1336def _collection_group_query_response_to_snapshot(
1337 response_pb: RunQueryResponse, collection
1338) -> Optional[document.DocumentSnapshot]:
1339 """Parse a query response protobuf to a document snapshot.
1340
1341 Args:
1342 response_pb (google.cloud.firestore_v1.\
1343 firestore.RunQueryResponse): A
1344 collection (:class:`~google.cloud.firestore_v1.collection.CollectionReference`):
1345 A reference to the collection that initiated the query.
1346
1347 Returns:
1348 Optional[:class:`~google.cloud.firestore.document.DocumentSnapshot`]:
1349 A snapshot of the data returned in the query. If
1350 ``response_pb.document`` is not set, the snapshot will be :data:`None`.
1351 """
1352 if not response_pb._pb.HasField("document"):
1353 return None
1354 reference = collection._client.document(response_pb.document.name)
1355 data = _helpers.decode_dict(response_pb.document.fields, collection._client)
1356 snapshot = document.DocumentSnapshot(
1357 reference,
1358 data,
1359 exists=True,
1360 read_time=response_pb._pb.read_time,
1361 create_time=response_pb._pb.document.create_time,
1362 update_time=response_pb._pb.document.update_time,
1363 )
1364 return snapshot
1365
1366
1367class BaseCollectionGroup(BaseQuery):
1368 """Represents a Collection Group in the Firestore API.
1369
1370 This is a specialization of :class:`.Query` that includes all documents in the
1371 database that are contained in a collection or subcollection of the given
1372 parent.
1373
1374 Args:
1375 parent (:class:`~google.cloud.firestore_v1.collection.CollectionReference`):
1376 The collection that this query applies to.
1377 """
1378
1379 _PARTITION_QUERY_ORDER = (
1380 BaseQuery._make_order(
1381 field_path_module.FieldPath.document_id(),
1382 BaseQuery.ASCENDING,
1383 ),
1384 )
1385
1386 def __init__(
1387 self,
1388 parent,
1389 projection=None,
1390 field_filters=(),
1391 orders=(),
1392 limit=None,
1393 limit_to_last=False,
1394 offset=None,
1395 start_at=None,
1396 end_at=None,
1397 all_descendants=True,
1398 recursive=False,
1399 ) -> None:
1400 if not all_descendants:
1401 raise ValueError("all_descendants must be True for collection group query.")
1402
1403 super(BaseCollectionGroup, self).__init__(
1404 parent=parent,
1405 projection=projection,
1406 field_filters=field_filters,
1407 orders=orders,
1408 limit=limit,
1409 limit_to_last=limit_to_last,
1410 offset=offset,
1411 start_at=start_at,
1412 end_at=end_at,
1413 all_descendants=all_descendants,
1414 recursive=recursive,
1415 )
1416
1417 def _validate_partition_query(self):
1418 if self._field_filters:
1419 raise ValueError("Can't partition query with filters.")
1420
1421 if self._projection:
1422 raise ValueError("Can't partition query with projection.")
1423
1424 if self._limit:
1425 raise ValueError("Can't partition query with limit.")
1426
1427 if self._offset:
1428 raise ValueError("Can't partition query with offset.")
1429
1430 def _get_query_class(self):
1431 raise NotImplementedError
1432
1433 def _prep_get_partitions(
1434 self,
1435 partition_count,
1436 retry: retries.Retry | object | None = None,
1437 timeout: float | None = None,
1438 read_time: datetime.datetime | None = None,
1439 ) -> Tuple[dict, dict]:
1440 self._validate_partition_query()
1441 parent_path, expected_prefix = self._parent._parent_info()
1442 klass = self._get_query_class()
1443 query = klass(
1444 self._parent,
1445 orders=self._PARTITION_QUERY_ORDER,
1446 start_at=self._start_at,
1447 end_at=self._end_at,
1448 all_descendants=self._all_descendants,
1449 )
1450 request = {
1451 "parent": parent_path,
1452 "structured_query": query._to_protobuf(),
1453 "partition_count": partition_count,
1454 }
1455 if read_time is not None:
1456 request["read_time"] = read_time
1457 kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
1458
1459 return request, kwargs
1460
1461 def get_partitions(
1462 self,
1463 partition_count,
1464 retry: Optional[retries.Retry] = None,
1465 timeout: Optional[float] = None,
1466 *,
1467 read_time: Optional[datetime.datetime] = None,
1468 ):
1469 raise NotImplementedError
1470
1471
1472class QueryPartition:
1473 """Represents a bounded partition of a collection group query.
1474
1475 Contains cursors that can be used in a query as a starting and/or end point for the
1476 collection group query. The cursors may only be used in a query that matches the
1477 constraints of the query that produced this partition.
1478
1479 Args:
1480 query (BaseQuery): The original query that this is a partition of.
1481 start_at (Optional[~google.cloud.firestore_v1.document.DocumentSnapshot]):
1482 Cursor for first query result to include. If `None`, the partition starts at
1483 the beginning of the result set.
1484 end_at (Optional[~google.cloud.firestore_v1.document.DocumentSnapshot]):
1485 Cursor for first query result after the last result included in the
1486 partition. If `None`, the partition runs to the end of the result set.
1487
1488 """
1489
1490 def __init__(self, query, start_at, end_at):
1491 self._query = query
1492 self._start_at = start_at
1493 self._end_at = end_at
1494
1495 @property
1496 def start_at(self):
1497 return self._start_at
1498
1499 @property
1500 def end_at(self):
1501 return self._end_at
1502
1503 def query(self):
1504 """Generate a new query using this partition's bounds.
1505
1506 Returns:
1507 BaseQuery: Copy of the original query with start and end bounds set by the
1508 cursors from this partition.
1509 """
1510 query = self._query
1511 start_at = ([self.start_at], True) if self.start_at else None
1512 end_at = ([self.end_at], True) if self.end_at else None
1513
1514 return type(query)(
1515 query._parent,
1516 all_descendants=query._all_descendants,
1517 orders=query._PARTITION_QUERY_ORDER,
1518 start_at=start_at,
1519 end_at=end_at,
1520 )