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