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 collections for the Google Cloud Firestore API."""
16from __future__ import annotations
17
18import random
19
20from typing import (
21 TYPE_CHECKING,
22 Any,
23 AsyncGenerator,
24 AsyncIterator,
25 Coroutine,
26 Generator,
27 Generic,
28 Iterable,
29 Sequence,
30 Tuple,
31 Union,
32 Optional,
33)
34
35from google.api_core import retry as retries
36
37from google.cloud.firestore_v1 import _helpers
38from google.cloud.firestore_v1.base_document import BaseDocumentReference
39from google.cloud.firestore_v1.base_query import QueryType
40
41if TYPE_CHECKING: # pragma: NO COVER
42 # Types needed only for Type Hints
43 from google.cloud.firestore_v1.base_aggregation import BaseAggregationQuery
44 from google.cloud.firestore_v1.base_document import DocumentSnapshot
45 from google.cloud.firestore_v1.base_vector_query import (
46 BaseVectorQuery,
47 DistanceMeasure,
48 )
49 from google.cloud.firestore_v1.async_document import AsyncDocumentReference
50 from google.cloud.firestore_v1.document import DocumentReference
51 from google.cloud.firestore_v1.field_path import FieldPath
52 from google.cloud.firestore_v1.query_profile import ExplainOptions
53 from google.cloud.firestore_v1.query_results import QueryResultsList
54 from google.cloud.firestore_v1.stream_generator import StreamGenerator
55 from google.cloud.firestore_v1.transaction import Transaction
56 from google.cloud.firestore_v1.vector import Vector
57 from google.cloud.firestore_v1.vector_query import VectorQuery
58
59 import datetime
60
61_AUTO_ID_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
62
63
64class BaseCollectionReference(Generic[QueryType]):
65 """A reference to a collection in a Firestore database.
66
67 The collection may already exist or this class can facilitate creation
68 of documents within the collection.
69
70 Args:
71 path (Tuple[str, ...]): The components in the collection path.
72 This is a series of strings representing each collection and
73 sub-collection ID, as well as the document IDs for any documents
74 that contain a sub-collection.
75 kwargs (dict): The keyword arguments for the constructor. The only
76 supported keyword is ``client`` and it must be a
77 :class:`~google.cloud.firestore_v1.client.Client` if provided. It
78 represents the client that created this collection reference.
79
80 Raises:
81 ValueError: if
82
83 * the ``path`` is empty
84 * there are an even number of elements
85 * a collection ID in ``path`` is not a string
86 * a document ID in ``path`` is not a string
87 TypeError: If a keyword other than ``client`` is used.
88 """
89
90 def __init__(self, *path, **kwargs) -> None:
91 _helpers.verify_path(path, is_collection=True)
92 self._path = path
93 self._client = kwargs.pop("client", None)
94 if kwargs:
95 raise TypeError(
96 "Received unexpected arguments", kwargs, "Only `client` is supported"
97 )
98
99 def __eq__(self, other):
100 if not isinstance(other, self.__class__):
101 return NotImplemented
102 return self._path == other._path and self._client == other._client
103
104 @property
105 def id(self):
106 """The collection identifier.
107
108 Returns:
109 str: The last component of the path.
110 """
111 return self._path[-1]
112
113 @property
114 def parent(self):
115 """Document that owns the current collection.
116
117 Returns:
118 Optional[:class:`~google.cloud.firestore_v1.document.DocumentReference`]:
119 The parent document, if the current collection is not a
120 top-level collection.
121 """
122 if len(self._path) == 1:
123 return None
124 else:
125 parent_path = self._path[:-1]
126 return self._client.document(*parent_path)
127
128 def _query(self) -> QueryType:
129 raise NotImplementedError
130
131 def _aggregation_query(self) -> BaseAggregationQuery:
132 raise NotImplementedError
133
134 def _vector_query(self) -> BaseVectorQuery:
135 raise NotImplementedError
136
137 def document(self, document_id: Optional[str] = None) -> BaseDocumentReference:
138 """Create a sub-document underneath the current collection.
139
140 Args:
141 document_id (Optional[str]): The document identifier
142 within the current collection. If not provided, will default
143 to a random 20 character string composed of digits,
144 uppercase and lowercase and letters.
145
146 Returns:
147 :class:`~google.cloud.firestore_v1.base_document.BaseDocumentReference`:
148 The child document.
149 """
150 if document_id is None:
151 document_id = _auto_id()
152
153 # Append `self._path` and the passed document's ID as long as the first
154 # element in the path is not an empty string, which comes from setting the
155 # parent to "" for recursive queries.
156 child_path = self._path + (document_id,) if self._path[0] else (document_id,)
157 return self._client.document(*child_path)
158
159 def _parent_info(self) -> Tuple[Any, str]:
160 """Get fully-qualified parent path and prefix for this collection.
161
162 Returns:
163 Tuple[str, str]: Pair of
164
165 * the fully-qualified (with database and project) path to the
166 parent of this collection (will either be the database path
167 or a document path).
168 * the prefix to a document in this collection.
169 """
170 parent_doc = self.parent
171 if parent_doc is None:
172 parent_path = _helpers.DOCUMENT_PATH_DELIMITER.join(
173 (self._client._database_string, "documents")
174 )
175 else:
176 parent_path = parent_doc._document_path
177
178 expected_prefix = _helpers.DOCUMENT_PATH_DELIMITER.join((parent_path, self.id))
179 return parent_path, expected_prefix
180
181 def _prep_add(
182 self,
183 document_data: dict,
184 document_id: Optional[str] = None,
185 retry: retries.Retry | retries.AsyncRetry | object | None = None,
186 timeout: Optional[float] = None,
187 ):
188 """Shared setup for async / sync :method:`add`"""
189 if document_id is None:
190 document_id = _auto_id()
191
192 document_ref = self.document(document_id)
193 kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
194
195 return document_ref, kwargs
196
197 def add(
198 self,
199 document_data: dict,
200 document_id: Optional[str] = None,
201 retry: retries.Retry | retries.AsyncRetry | object | None = None,
202 timeout: Optional[float] = None,
203 ) -> Union[Tuple[Any, Any], Coroutine[Any, Any, Tuple[Any, Any]]]:
204 raise NotImplementedError
205
206 def _prep_list_documents(
207 self,
208 page_size: Optional[int] = None,
209 retry: retries.Retry | retries.AsyncRetry | object | None = None,
210 timeout: Optional[float] = None,
211 read_time: Optional[datetime.datetime] = None,
212 ) -> Tuple[dict, dict]:
213 """Shared setup for async / sync :method:`list_documents`"""
214 parent, _ = self._parent_info()
215 request = {
216 "parent": parent,
217 "collection_id": self.id,
218 "page_size": page_size,
219 "show_missing": True,
220 # list_documents returns an iterator of document references, which do not
221 # include any fields. To save on data transfer, we can set a field_path mask
222 # to include no fields
223 "mask": {"field_paths": None},
224 }
225 if read_time is not None:
226 request["read_time"] = read_time
227 kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
228
229 return request, kwargs
230
231 def list_documents(
232 self,
233 page_size: Optional[int] = None,
234 retry: retries.Retry | retries.AsyncRetry | object | None = None,
235 timeout: Optional[float] = None,
236 *,
237 read_time: Optional[datetime.datetime] = None,
238 ) -> Union[
239 Generator[DocumentReference, Any, Any],
240 AsyncGenerator[AsyncDocumentReference, Any],
241 ]:
242 raise NotImplementedError
243
244 def recursive(self) -> QueryType:
245 return self._query().recursive()
246
247 def select(self, field_paths: Iterable[str]) -> QueryType:
248 """Create a "select" query with this collection as parent.
249
250 See
251 :meth:`~google.cloud.firestore_v1.query.Query.select` for
252 more information on this method.
253
254 Args:
255 field_paths (Iterable[str, ...]): An iterable of field paths
256 (``.``-delimited list of field names) to use as a projection
257 of document fields in the query results.
258
259 Returns:
260 :class:`~google.cloud.firestore_v1.query.Query`:
261 A "projected" query.
262 """
263 query = self._query()
264 return query.select(field_paths)
265
266 def where(
267 self,
268 field_path: Optional[str] = None,
269 op_string: Optional[str] = None,
270 value=None,
271 *,
272 filter=None,
273 ) -> QueryType:
274 """Create a "where" query with this collection as parent.
275
276 See
277 :meth:`~google.cloud.firestore_v1.query.Query.where` for
278 more information on this method.
279
280 Args:
281 field_path (str): A field path (``.``-delimited list of
282 field names) for the field to filter on. Optional.
283 op_string (str): A comparison operation in the form of a string.
284 Acceptable values are ``<``, ``<=``, ``==``, ``>=``, ``>``,
285 and ``in``. Optional.
286 value (Any): The value to compare the field against in the filter.
287 If ``value`` is :data:`None` or a NaN, then ``==`` is the only
288 allowed operation. If ``op_string`` is ``in``, ``value``
289 must be a sequence of values. Optional.
290 filter (class:`~google.cloud.firestore_v1.base_query.BaseFilter`): an instance of a Filter.
291 Either a FieldFilter or a CompositeFilter.
292 Returns:
293 :class:`~google.cloud.firestore_v1.query.Query`:
294 A filtered query.
295 Raises:
296 ValueError, if both the positional arguments (field_path, op_string, value)
297 and the filter keyword argument are passed at the same time.
298 """
299 query = self._query()
300 if field_path and op_string:
301 if filter is not None:
302 raise ValueError(
303 "Can't pass in both the positional arguments and 'filter' at the same time"
304 )
305 if field_path == "__name__" and op_string == "in":
306 wrapped_names = []
307
308 for name in value:
309 if isinstance(name, str):
310 name = self.document(name)
311
312 wrapped_names.append(name)
313
314 value = wrapped_names
315 return query.where(field_path, op_string, value)
316 else:
317 return query.where(filter=filter)
318
319 def order_by(self, field_path: str, **kwargs) -> QueryType:
320 """Create an "order by" query with this collection as parent.
321
322 See
323 :meth:`~google.cloud.firestore_v1.query.Query.order_by` for
324 more information on this method.
325
326 Args:
327 field_path (str): A field path (``.``-delimited list of
328 field names) on which to order the query results.
329 kwargs (Dict[str, Any]): The keyword arguments to pass along
330 to the query. The only supported keyword is ``direction``,
331 see :meth:`~google.cloud.firestore_v1.query.Query.order_by`
332 for more information.
333
334 Returns:
335 :class:`~google.cloud.firestore_v1.query.Query`:
336 An "order by" query.
337 """
338 query = self._query()
339 return query.order_by(field_path, **kwargs)
340
341 def limit(self, count: int) -> QueryType:
342 """Create a limited query with this collection as parent.
343
344 .. note::
345 `limit` and `limit_to_last` are mutually exclusive.
346 Setting `limit` will drop previously set `limit_to_last`.
347
348 See
349 :meth:`~google.cloud.firestore_v1.query.Query.limit` for
350 more information on this method.
351
352 Args:
353 count (int): Maximum number of documents to return that match
354 the query.
355
356 Returns:
357 :class:`~google.cloud.firestore_v1.query.Query`:
358 A limited query.
359 """
360 query = self._query()
361 return query.limit(count)
362
363 def limit_to_last(self, count: int):
364 """Create a limited to last query with this collection as parent.
365
366 .. note::
367 `limit` and `limit_to_last` are mutually exclusive.
368 Setting `limit_to_last` will drop previously set `limit`.
369
370 See
371 :meth:`~google.cloud.firestore_v1.query.Query.limit_to_last`
372 for more information on this method.
373
374 Args:
375 count (int): Maximum number of documents to return that
376 match the query.
377 Returns:
378 :class:`~google.cloud.firestore_v1.query.Query`:
379 A limited to last query.
380 """
381 query = self._query()
382 return query.limit_to_last(count)
383
384 def offset(self, num_to_skip: int) -> QueryType:
385 """Skip to an offset in a query with this collection as parent.
386
387 See
388 :meth:`~google.cloud.firestore_v1.query.Query.offset` for
389 more information on this method.
390
391 Args:
392 num_to_skip (int): The number of results to skip at the beginning
393 of query results. (Must be non-negative.)
394
395 Returns:
396 :class:`~google.cloud.firestore_v1.query.Query`:
397 An offset query.
398 """
399 query = self._query()
400 return query.offset(num_to_skip)
401
402 def start_at(
403 self, document_fields: Union[DocumentSnapshot, dict, list, tuple]
404 ) -> QueryType:
405 """Start query at a cursor with this collection as parent.
406
407 See
408 :meth:`~google.cloud.firestore_v1.query.Query.start_at` for
409 more information on this method.
410
411 Args:
412 document_fields (Union[:class:`~google.cloud.firestore_v1.\
413 document.DocumentSnapshot`, dict, list, tuple]):
414 A document snapshot or a dictionary/list/tuple of fields
415 representing a query results cursor. A cursor is a collection
416 of values that represent a position in a query result set.
417
418 Returns:
419 :class:`~google.cloud.firestore_v1.query.Query`:
420 A query with cursor.
421 """
422 query = self._query()
423 return query.start_at(document_fields)
424
425 def start_after(
426 self, document_fields: Union[DocumentSnapshot, dict, list, tuple]
427 ) -> QueryType:
428 """Start query after a cursor with this collection as parent.
429
430 See
431 :meth:`~google.cloud.firestore_v1.query.Query.start_after` for
432 more information on this method.
433
434 Args:
435 document_fields (Union[:class:`~google.cloud.firestore_v1.\
436 document.DocumentSnapshot`, dict, list, tuple]):
437 A document snapshot or a dictionary/list/tuple of fields
438 representing a query results cursor. A cursor is a collection
439 of values that represent a position in a query result set.
440
441 Returns:
442 :class:`~google.cloud.firestore_v1.query.Query`:
443 A query with cursor.
444 """
445 query = self._query()
446 return query.start_after(document_fields)
447
448 def end_before(
449 self, document_fields: Union[DocumentSnapshot, dict, list, tuple]
450 ) -> QueryType:
451 """End query before a cursor with this collection as parent.
452
453 See
454 :meth:`~google.cloud.firestore_v1.query.Query.end_before` for
455 more information on this method.
456
457 Args:
458 document_fields (Union[:class:`~google.cloud.firestore_v1.\
459 document.DocumentSnapshot`, dict, list, tuple]):
460 A document snapshot or a dictionary/list/tuple of fields
461 representing a query results cursor. A cursor is a collection
462 of values that represent a position in a query result set.
463
464 Returns:
465 :class:`~google.cloud.firestore_v1.query.Query`:
466 A query with cursor.
467 """
468 query = self._query()
469 return query.end_before(document_fields)
470
471 def end_at(
472 self, document_fields: Union[DocumentSnapshot, dict, list, tuple]
473 ) -> QueryType:
474 """End query at a cursor with this collection as parent.
475
476 See
477 :meth:`~google.cloud.firestore_v1.query.Query.end_at` for
478 more information on this method.
479
480 Args:
481 document_fields (Union[:class:`~google.cloud.firestore_v1.\
482 document.DocumentSnapshot`, dict, list, tuple]):
483 A document snapshot or a dictionary/list/tuple of fields
484 representing a query results cursor. A cursor is a collection
485 of values that represent a position in a query result set.
486
487 Returns:
488 :class:`~google.cloud.firestore_v1.query.Query`:
489 A query with cursor.
490 """
491 query = self._query()
492 return query.end_at(document_fields)
493
494 def _prep_get_or_stream(
495 self,
496 retry: retries.Retry | retries.AsyncRetry | object | None = None,
497 timeout: Optional[float] = None,
498 ) -> Tuple[Any, dict]:
499 """Shared setup for async / sync :meth:`get` / :meth:`stream`"""
500 query = self._query()
501 kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
502
503 return query, kwargs
504
505 def get(
506 self,
507 transaction: Optional[Transaction] = None,
508 retry: retries.Retry | retries.AsyncRetry | object | None = None,
509 timeout: Optional[float] = None,
510 *,
511 explain_options: Optional[ExplainOptions] = None,
512 read_time: Optional[datetime.datetime] = None,
513 ) -> (
514 QueryResultsList[DocumentSnapshot]
515 | Coroutine[Any, Any, QueryResultsList[DocumentSnapshot]]
516 ):
517 raise NotImplementedError
518
519 def stream(
520 self,
521 transaction: Optional[Transaction] = None,
522 retry: retries.Retry | retries.AsyncRetry | object | None = None,
523 timeout: Optional[float] = None,
524 *,
525 explain_options: Optional[ExplainOptions] = None,
526 read_time: Optional[datetime.datetime] = None,
527 ) -> StreamGenerator[DocumentSnapshot] | AsyncIterator[DocumentSnapshot]:
528 raise NotImplementedError
529
530 def on_snapshot(self, callback):
531 raise NotImplementedError
532
533 def count(self, alias=None):
534 """
535 Adds a count over the nested query.
536
537 :type alias: str
538 :param alias: (Optional) The alias for the count
539 """
540 return self._aggregation_query().count(alias=alias)
541
542 def sum(self, field_ref: str | FieldPath, alias=None):
543 """
544 Adds a sum over the nested query.
545
546 :type field_ref: Union[str, google.cloud.firestore_v1.field_path.FieldPath]
547 :param field_ref: The field to aggregate across.
548
549 :type alias: Optional[str]
550 :param alias: Optional name of the field to store the result of the aggregation into.
551 If not provided, Firestore will pick a default name following the format field_<incremental_id++>.
552
553 """
554 return self._aggregation_query().sum(field_ref, alias=alias)
555
556 def avg(self, field_ref: str | FieldPath, alias=None):
557 """
558 Adds an avg over the nested query.
559
560 :type field_ref: Union[str, google.cloud.firestore_v1.field_path.FieldPath]
561 :param field_ref: The field to aggregate across.
562
563 :type alias: Optional[str]
564 :param alias: Optional name of the field to store the result of the aggregation into.
565 If not provided, Firestore will pick a default name following the format field_<incremental_id++>.
566 """
567 return self._aggregation_query().avg(field_ref, alias=alias)
568
569 def find_nearest(
570 self,
571 vector_field: str,
572 query_vector: Union[Vector, Sequence[float]],
573 limit: int,
574 distance_measure: DistanceMeasure,
575 *,
576 distance_result_field: Optional[str] = None,
577 distance_threshold: Optional[float] = None,
578 ) -> VectorQuery:
579 """
580 Finds the closest vector embeddings to the given query vector.
581
582 Args:
583 vector_field (str): An indexed vector field to search upon. Only documents which contain
584 vectors whose dimensionality match the query_vector can be returned.
585 query_vector(Union[Vector, Sequence[float]]): The query vector that we are searching on. Must be a vector of no more
586 than 2048 dimensions.
587 limit (int): The number of nearest neighbors to return. Must be a positive integer of no more than 1000.
588 distance_measure (:class:`DistanceMeasure`): The Distance Measure to use.
589 distance_result_field (Optional[str]):
590 Name of the field to output the result of the vector distance calculation
591 distance_threshold (Optional[float]):
592 A threshold for which no less similar documents will be returned.
593
594 Returns:
595 :class`~firestore_v1.vector_query.VectorQuery`: the vector query.
596 """
597 return self._vector_query().find_nearest(
598 vector_field,
599 query_vector,
600 limit,
601 distance_measure,
602 distance_result_field=distance_result_field,
603 distance_threshold=distance_threshold,
604 )
605
606
607def _auto_id() -> str:
608 """Generate a "random" automatically generated ID.
609
610 Returns:
611 str: A 20 character string composed of digits, uppercase and
612 lowercase and letters.
613 """
614
615 return "".join(random.choice(_AUTO_ID_CHARS) for _ in range(20))
616
617
618def _item_to_document_ref(collection_reference, item):
619 """Convert Document resource to document ref.
620
621 Args:
622 collection_reference (google.api_core.page_iterator.GRPCIterator):
623 iterator response
624 item (dict): document resource
625
626 Returns:
627 :class:`~google.cloud.firestore_v1.base_document.BaseDocumentReference`:
628 The child document
629 """
630 document_id = item.name.split(_helpers.DOCUMENT_PATH_DELIMITER)[-1]
631 return collection_reference.document(document_id)