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

130 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_query import QueryType 

39 

40if TYPE_CHECKING: # pragma: NO COVER 

41 # Types needed only for Type Hints 

42 from google.cloud.firestore_v1.base_aggregation import BaseAggregationQuery 

43 from google.cloud.firestore_v1.base_document import DocumentSnapshot 

44 from google.cloud.firestore_v1.base_vector_query import ( 

45 BaseVectorQuery, 

46 DistanceMeasure, 

47 ) 

48 from google.cloud.firestore_v1.async_document import AsyncDocumentReference 

49 from google.cloud.firestore_v1.document import DocumentReference 

50 from google.cloud.firestore_v1.field_path import FieldPath 

51 from google.cloud.firestore_v1.query_profile import ExplainOptions 

52 from google.cloud.firestore_v1.query_results import QueryResultsList 

53 from google.cloud.firestore_v1.stream_generator import StreamGenerator 

54 from google.cloud.firestore_v1.transaction import Transaction 

55 from google.cloud.firestore_v1.vector import Vector 

56 from google.cloud.firestore_v1.vector_query import VectorQuery 

57 

58 import datetime 

59 

60_AUTO_ID_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" 

61 

62 

63class BaseCollectionReference(Generic[QueryType]): 

64 """A reference to a collection in a Firestore database. 

65 

66 The collection may already exist or this class can facilitate creation 

67 of documents within the collection. 

68 

69 Args: 

70 path (Tuple[str, ...]): The components in the collection path. 

71 This is a series of strings representing each collection and 

72 sub-collection ID, as well as the document IDs for any documents 

73 that contain a sub-collection. 

74 kwargs (dict): The keyword arguments for the constructor. The only 

75 supported keyword is ``client`` and it must be a 

76 :class:`~google.cloud.firestore_v1.client.Client` if provided. It 

77 represents the client that created this collection reference. 

78 

79 Raises: 

80 ValueError: if 

81 

82 * the ``path`` is empty 

83 * there are an even number of elements 

84 * a collection ID in ``path`` is not a string 

85 * a document ID in ``path`` is not a string 

86 TypeError: If a keyword other than ``client`` is used. 

87 """ 

88 

89 def __init__(self, *path, **kwargs) -> None: 

90 _helpers.verify_path(path, is_collection=True) 

91 self._path = path 

92 self._client = kwargs.pop("client", None) 

93 if kwargs: 

94 raise TypeError( 

95 "Received unexpected arguments", kwargs, "Only `client` is supported" 

96 ) 

97 

98 def __eq__(self, other): 

99 if not isinstance(other, self.__class__): 

100 return NotImplemented 

101 return self._path == other._path and self._client == other._client 

102 

103 @property 

104 def id(self): 

105 """The collection identifier. 

106 

107 Returns: 

108 str: The last component of the path. 

109 """ 

110 return self._path[-1] 

111 

112 @property 

113 def parent(self): 

114 """Document that owns the current collection. 

115 

116 Returns: 

117 Optional[:class:`~google.cloud.firestore_v1.document.DocumentReference`]: 

118 The parent document, if the current collection is not a 

119 top-level collection. 

120 """ 

121 if len(self._path) == 1: 

122 return None 

123 else: 

124 parent_path = self._path[:-1] 

125 return self._client.document(*parent_path) 

126 

127 def _query(self) -> QueryType: 

128 raise NotImplementedError 

129 

130 def _aggregation_query(self) -> BaseAggregationQuery: 

131 raise NotImplementedError 

132 

133 def _vector_query(self) -> BaseVectorQuery: 

134 raise NotImplementedError 

135 

136 def document(self, document_id: Optional[str] = None): 

137 """Create a sub-document underneath the current collection. 

138 

139 Args: 

140 document_id (Optional[str]): The document identifier 

141 within the current collection. If not provided, will default 

142 to a random 20 character string composed of digits, 

143 uppercase and lowercase and letters. 

144 

145 Returns: 

146 :class:`~google.cloud.firestore_v1.base_document.BaseDocumentReference`: 

147 The child document. 

148 """ 

149 if document_id is None: 

150 document_id = _auto_id() 

151 

152 # Append `self._path` and the passed document's ID as long as the first 

153 # element in the path is not an empty string, which comes from setting the 

154 # parent to "" for recursive queries. 

155 child_path = self._path + (document_id,) if self._path[0] else (document_id,) 

156 return self._client.document(*child_path) 

157 

158 def _parent_info(self) -> Tuple[Any, str]: 

159 """Get fully-qualified parent path and prefix for this collection. 

160 

161 Returns: 

162 Tuple[str, str]: Pair of 

163 

164 * the fully-qualified (with database and project) path to the 

165 parent of this collection (will either be the database path 

166 or a document path). 

167 * the prefix to a document in this collection. 

168 """ 

169 parent_doc = self.parent 

170 if parent_doc is None: 

171 parent_path = _helpers.DOCUMENT_PATH_DELIMITER.join( 

172 (self._client._database_string, "documents") 

173 ) 

174 else: 

175 parent_path = parent_doc._document_path 

176 

177 expected_prefix = _helpers.DOCUMENT_PATH_DELIMITER.join((parent_path, self.id)) 

178 return parent_path, expected_prefix 

179 

180 def _prep_add( 

181 self, 

182 document_data: dict, 

183 document_id: Optional[str] = None, 

184 retry: retries.Retry | retries.AsyncRetry | object | None = None, 

185 timeout: Optional[float] = None, 

186 ): 

187 """Shared setup for async / sync :method:`add`""" 

188 if document_id is None: 

189 document_id = _auto_id() 

190 

191 document_ref = self.document(document_id) 

192 kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) 

193 

194 return document_ref, kwargs 

195 

196 def add( 

197 self, 

198 document_data: dict, 

199 document_id: Optional[str] = None, 

200 retry: retries.Retry | retries.AsyncRetry | object | None = None, 

201 timeout: Optional[float] = None, 

202 ) -> Union[Tuple[Any, Any], Coroutine[Any, Any, Tuple[Any, Any]]]: 

203 raise NotImplementedError 

204 

205 def _prep_list_documents( 

206 self, 

207 page_size: Optional[int] = None, 

208 retry: retries.Retry | retries.AsyncRetry | object | None = None, 

209 timeout: Optional[float] = None, 

210 read_time: Optional[datetime.datetime] = None, 

211 ) -> Tuple[dict, dict]: 

212 """Shared setup for async / sync :method:`list_documents`""" 

213 parent, _ = self._parent_info() 

214 request = { 

215 "parent": parent, 

216 "collection_id": self.id, 

217 "page_size": page_size, 

218 "show_missing": True, 

219 # list_documents returns an iterator of document references, which do not 

220 # include any fields. To save on data transfer, we can set a field_path mask 

221 # to include no fields 

222 "mask": {"field_paths": None}, 

223 } 

224 if read_time is not None: 

225 request["read_time"] = read_time 

226 kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) 

227 

228 return request, kwargs 

229 

230 def list_documents( 

231 self, 

232 page_size: Optional[int] = None, 

233 retry: retries.Retry | retries.AsyncRetry | object | None = None, 

234 timeout: Optional[float] = None, 

235 *, 

236 read_time: Optional[datetime.datetime] = None, 

237 ) -> Union[ 

238 Generator[DocumentReference, Any, Any], 

239 AsyncGenerator[AsyncDocumentReference, Any], 

240 ]: 

241 raise NotImplementedError 

242 

243 def recursive(self) -> QueryType: 

244 return self._query().recursive() 

245 

246 def select(self, field_paths: Iterable[str]) -> QueryType: 

247 """Create a "select" query with this collection as parent. 

248 

249 See 

250 :meth:`~google.cloud.firestore_v1.query.Query.select` for 

251 more information on this method. 

252 

253 Args: 

254 field_paths (Iterable[str, ...]): An iterable of field paths 

255 (``.``-delimited list of field names) to use as a projection 

256 of document fields in the query results. 

257 

258 Returns: 

259 :class:`~google.cloud.firestore_v1.query.Query`: 

260 A "projected" query. 

261 """ 

262 query = self._query() 

263 return query.select(field_paths) 

264 

265 def where( 

266 self, 

267 field_path: Optional[str] = None, 

268 op_string: Optional[str] = None, 

269 value=None, 

270 *, 

271 filter=None, 

272 ) -> QueryType: 

273 """Create a "where" query with this collection as parent. 

274 

275 See 

276 :meth:`~google.cloud.firestore_v1.query.Query.where` for 

277 more information on this method. 

278 

279 Args: 

280 field_path (str): A field path (``.``-delimited list of 

281 field names) for the field to filter on. Optional. 

282 op_string (str): A comparison operation in the form of a string. 

283 Acceptable values are ``<``, ``<=``, ``==``, ``>=``, ``>``, 

284 and ``in``. Optional. 

285 value (Any): The value to compare the field against in the filter. 

286 If ``value`` is :data:`None` or a NaN, then ``==`` is the only 

287 allowed operation. If ``op_string`` is ``in``, ``value`` 

288 must be a sequence of values. Optional. 

289 filter (class:`~google.cloud.firestore_v1.base_query.BaseFilter`): an instance of a Filter. 

290 Either a FieldFilter or a CompositeFilter. 

291 Returns: 

292 :class:`~google.cloud.firestore_v1.query.Query`: 

293 A filtered query. 

294 Raises: 

295 ValueError, if both the positional arguments (field_path, op_string, value) 

296 and the filter keyword argument are passed at the same time. 

297 """ 

298 query = self._query() 

299 if field_path and op_string: 

300 if filter is not None: 

301 raise ValueError( 

302 "Can't pass in both the positional arguments and 'filter' at the same time" 

303 ) 

304 if field_path == "__name__" and op_string == "in": 

305 wrapped_names = [] 

306 

307 for name in value: 

308 if isinstance(name, str): 

309 name = self.document(name) 

310 

311 wrapped_names.append(name) 

312 

313 value = wrapped_names 

314 return query.where(field_path, op_string, value) 

315 else: 

316 return query.where(filter=filter) 

317 

318 def order_by(self, field_path: str, **kwargs) -> QueryType: 

319 """Create an "order by" query with this collection as parent. 

320 

321 See 

322 :meth:`~google.cloud.firestore_v1.query.Query.order_by` for 

323 more information on this method. 

324 

325 Args: 

326 field_path (str): A field path (``.``-delimited list of 

327 field names) on which to order the query results. 

328 kwargs (Dict[str, Any]): The keyword arguments to pass along 

329 to the query. The only supported keyword is ``direction``, 

330 see :meth:`~google.cloud.firestore_v1.query.Query.order_by` 

331 for more information. 

332 

333 Returns: 

334 :class:`~google.cloud.firestore_v1.query.Query`: 

335 An "order by" query. 

336 """ 

337 query = self._query() 

338 return query.order_by(field_path, **kwargs) 

339 

340 def limit(self, count: int) -> QueryType: 

341 """Create a limited query with this collection as parent. 

342 

343 .. note:: 

344 `limit` and `limit_to_last` are mutually exclusive. 

345 Setting `limit` will drop previously set `limit_to_last`. 

346 

347 See 

348 :meth:`~google.cloud.firestore_v1.query.Query.limit` for 

349 more information on this method. 

350 

351 Args: 

352 count (int): Maximum number of documents to return that match 

353 the query. 

354 

355 Returns: 

356 :class:`~google.cloud.firestore_v1.query.Query`: 

357 A limited query. 

358 """ 

359 query = self._query() 

360 return query.limit(count) 

361 

362 def limit_to_last(self, count: int): 

363 """Create a limited to last query with this collection as parent. 

364 

365 .. note:: 

366 `limit` and `limit_to_last` are mutually exclusive. 

367 Setting `limit_to_last` will drop previously set `limit`. 

368 

369 See 

370 :meth:`~google.cloud.firestore_v1.query.Query.limit_to_last` 

371 for more information on this method. 

372 

373 Args: 

374 count (int): Maximum number of documents to return that 

375 match the query. 

376 Returns: 

377 :class:`~google.cloud.firestore_v1.query.Query`: 

378 A limited to last query. 

379 """ 

380 query = self._query() 

381 return query.limit_to_last(count) 

382 

383 def offset(self, num_to_skip: int) -> QueryType: 

384 """Skip to an offset in a query with this collection as parent. 

385 

386 See 

387 :meth:`~google.cloud.firestore_v1.query.Query.offset` for 

388 more information on this method. 

389 

390 Args: 

391 num_to_skip (int): The number of results to skip at the beginning 

392 of query results. (Must be non-negative.) 

393 

394 Returns: 

395 :class:`~google.cloud.firestore_v1.query.Query`: 

396 An offset query. 

397 """ 

398 query = self._query() 

399 return query.offset(num_to_skip) 

400 

401 def start_at( 

402 self, document_fields: Union[DocumentSnapshot, dict, list, tuple] 

403 ) -> QueryType: 

404 """Start query at a cursor with this collection as parent. 

405 

406 See 

407 :meth:`~google.cloud.firestore_v1.query.Query.start_at` for 

408 more information on this method. 

409 

410 Args: 

411 document_fields (Union[:class:`~google.cloud.firestore_v1.\ 

412 document.DocumentSnapshot`, dict, list, tuple]): 

413 A document snapshot or a dictionary/list/tuple of fields 

414 representing a query results cursor. A cursor is a collection 

415 of values that represent a position in a query result set. 

416 

417 Returns: 

418 :class:`~google.cloud.firestore_v1.query.Query`: 

419 A query with cursor. 

420 """ 

421 query = self._query() 

422 return query.start_at(document_fields) 

423 

424 def start_after( 

425 self, document_fields: Union[DocumentSnapshot, dict, list, tuple] 

426 ) -> QueryType: 

427 """Start query after a cursor with this collection as parent. 

428 

429 See 

430 :meth:`~google.cloud.firestore_v1.query.Query.start_after` for 

431 more information on this method. 

432 

433 Args: 

434 document_fields (Union[:class:`~google.cloud.firestore_v1.\ 

435 document.DocumentSnapshot`, dict, list, tuple]): 

436 A document snapshot or a dictionary/list/tuple of fields 

437 representing a query results cursor. A cursor is a collection 

438 of values that represent a position in a query result set. 

439 

440 Returns: 

441 :class:`~google.cloud.firestore_v1.query.Query`: 

442 A query with cursor. 

443 """ 

444 query = self._query() 

445 return query.start_after(document_fields) 

446 

447 def end_before( 

448 self, document_fields: Union[DocumentSnapshot, dict, list, tuple] 

449 ) -> QueryType: 

450 """End query before a cursor with this collection as parent. 

451 

452 See 

453 :meth:`~google.cloud.firestore_v1.query.Query.end_before` for 

454 more information on this method. 

455 

456 Args: 

457 document_fields (Union[:class:`~google.cloud.firestore_v1.\ 

458 document.DocumentSnapshot`, dict, list, tuple]): 

459 A document snapshot or a dictionary/list/tuple of fields 

460 representing a query results cursor. A cursor is a collection 

461 of values that represent a position in a query result set. 

462 

463 Returns: 

464 :class:`~google.cloud.firestore_v1.query.Query`: 

465 A query with cursor. 

466 """ 

467 query = self._query() 

468 return query.end_before(document_fields) 

469 

470 def end_at( 

471 self, document_fields: Union[DocumentSnapshot, dict, list, tuple] 

472 ) -> QueryType: 

473 """End query at a cursor with this collection as parent. 

474 

475 See 

476 :meth:`~google.cloud.firestore_v1.query.Query.end_at` for 

477 more information on this method. 

478 

479 Args: 

480 document_fields (Union[:class:`~google.cloud.firestore_v1.\ 

481 document.DocumentSnapshot`, dict, list, tuple]): 

482 A document snapshot or a dictionary/list/tuple of fields 

483 representing a query results cursor. A cursor is a collection 

484 of values that represent a position in a query result set. 

485 

486 Returns: 

487 :class:`~google.cloud.firestore_v1.query.Query`: 

488 A query with cursor. 

489 """ 

490 query = self._query() 

491 return query.end_at(document_fields) 

492 

493 def _prep_get_or_stream( 

494 self, 

495 retry: retries.Retry | retries.AsyncRetry | object | None = None, 

496 timeout: Optional[float] = None, 

497 ) -> Tuple[Any, dict]: 

498 """Shared setup for async / sync :meth:`get` / :meth:`stream`""" 

499 query = self._query() 

500 kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) 

501 

502 return query, kwargs 

503 

504 def get( 

505 self, 

506 transaction: Optional[Transaction] = None, 

507 retry: retries.Retry | retries.AsyncRetry | object | None = None, 

508 timeout: Optional[float] = None, 

509 *, 

510 explain_options: Optional[ExplainOptions] = None, 

511 read_time: Optional[datetime.datetime] = None, 

512 ) -> ( 

513 QueryResultsList[DocumentSnapshot] 

514 | Coroutine[Any, Any, QueryResultsList[DocumentSnapshot]] 

515 ): 

516 raise NotImplementedError 

517 

518 def stream( 

519 self, 

520 transaction: Optional[Transaction] = None, 

521 retry: retries.Retry | retries.AsyncRetry | object | None = None, 

522 timeout: Optional[float] = None, 

523 *, 

524 explain_options: Optional[ExplainOptions] = None, 

525 read_time: Optional[datetime.datetime] = None, 

526 ) -> StreamGenerator[DocumentSnapshot] | AsyncIterator[DocumentSnapshot]: 

527 raise NotImplementedError 

528 

529 def on_snapshot(self, callback): 

530 raise NotImplementedError 

531 

532 def count(self, alias=None): 

533 """ 

534 Adds a count over the nested query. 

535 

536 :type alias: str 

537 :param alias: (Optional) The alias for the count 

538 """ 

539 return self._aggregation_query().count(alias=alias) 

540 

541 def sum(self, field_ref: str | FieldPath, alias=None): 

542 """ 

543 Adds a sum over the nested query. 

544 

545 :type field_ref: Union[str, google.cloud.firestore_v1.field_path.FieldPath] 

546 :param field_ref: The field to aggregate across. 

547 

548 :type alias: Optional[str] 

549 :param alias: Optional name of the field to store the result of the aggregation into. 

550 If not provided, Firestore will pick a default name following the format field_<incremental_id++>. 

551 

552 """ 

553 return self._aggregation_query().sum(field_ref, alias=alias) 

554 

555 def avg(self, field_ref: str | FieldPath, alias=None): 

556 """ 

557 Adds an avg over the nested query. 

558 

559 :type field_ref: Union[str, google.cloud.firestore_v1.field_path.FieldPath] 

560 :param field_ref: The field to aggregate across. 

561 

562 :type alias: Optional[str] 

563 :param alias: Optional name of the field to store the result of the aggregation into. 

564 If not provided, Firestore will pick a default name following the format field_<incremental_id++>. 

565 """ 

566 return self._aggregation_query().avg(field_ref, alias=alias) 

567 

568 def find_nearest( 

569 self, 

570 vector_field: str, 

571 query_vector: Union[Vector, Sequence[float]], 

572 limit: int, 

573 distance_measure: DistanceMeasure, 

574 *, 

575 distance_result_field: Optional[str] = None, 

576 distance_threshold: Optional[float] = None, 

577 ) -> VectorQuery: 

578 """ 

579 Finds the closest vector embeddings to the given query vector. 

580 

581 Args: 

582 vector_field (str): An indexed vector field to search upon. Only documents which contain 

583 vectors whose dimensionality match the query_vector can be returned. 

584 query_vector(Union[Vector, Sequence[float]]): The query vector that we are searching on. Must be a vector of no more 

585 than 2048 dimensions. 

586 limit (int): The number of nearest neighbors to return. Must be a positive integer of no more than 1000. 

587 distance_measure (:class:`DistanceMeasure`): The Distance Measure to use. 

588 distance_result_field (Optional[str]): 

589 Name of the field to output the result of the vector distance calculation 

590 distance_threshold (Optional[float]): 

591 A threshold for which no less similar documents will be returned. 

592 

593 Returns: 

594 :class`~firestore_v1.vector_query.VectorQuery`: the vector query. 

595 """ 

596 return self._vector_query().find_nearest( 

597 vector_field, 

598 query_vector, 

599 limit, 

600 distance_measure, 

601 distance_result_field=distance_result_field, 

602 distance_threshold=distance_threshold, 

603 ) 

604 

605 

606def _auto_id() -> str: 

607 """Generate a "random" automatically generated ID. 

608 

609 Returns: 

610 str: A 20 character string composed of digits, uppercase and 

611 lowercase and letters. 

612 """ 

613 

614 return "".join(random.choice(_AUTO_ID_CHARS) for _ in range(20)) 

615 

616 

617def _item_to_document_ref(collection_reference, item): 

618 """Convert Document resource to document ref. 

619 

620 Args: 

621 collection_reference (google.api_core.page_iterator.GRPCIterator): 

622 iterator response 

623 item (dict): document resource 

624 

625 Returns: 

626 :class:`~google.cloud.firestore_v1.base_document.BaseDocumentReference`: 

627 The child document 

628 """ 

629 document_id = item.name.split(_helpers.DOCUMENT_PATH_DELIMITER)[-1] 

630 return collection_reference.document(document_id)