Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/google/cloud/firestore_v1/base_collection.py: 38%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

133 statements  

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)