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

131 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.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)