1# Copyright 2021 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"""Helpers for batch requests to the Google Cloud Firestore API."""
16from __future__ import annotations
17from google.api_core import gapic_v1
18from google.api_core import retry as retries
19
20from google.cloud.firestore_v1 import _helpers
21from google.cloud.firestore_v1.base_batch import BaseBatch
22from google.cloud.firestore_v1.types.firestore import BatchWriteResponse
23
24
25class BulkWriteBatch(BaseBatch):
26 """Accumulate write operations to be sent in a batch. Use this over
27 `WriteBatch` for higher volumes (e.g., via `BulkWriter`) and when the order
28 of operations within a given batch is unimportant.
29
30 Because the order in which individual write operations are applied to the database
31 is not guaranteed, `batch_write` RPCs can never contain multiple operations
32 to the same document. If calling code detects a second write operation to a
33 known document reference, it should first cut off the previous batch and
34 send it, then create a new batch starting with the latest write operation.
35 In practice, the [Async]BulkWriter classes handle this.
36
37 This has the same set of methods for write operations that
38 :class:`~google.cloud.firestore_v1.document.DocumentReference` does,
39 e.g. :meth:`~google.cloud.firestore_v1.document.DocumentReference.create`.
40
41 Args:
42 client (:class:`~google.cloud.firestore_v1.client.Client`):
43 The client that created this batch.
44 """
45
46 def __init__(self, client) -> None:
47 super(BulkWriteBatch, self).__init__(client=client)
48
49 def commit(
50 self,
51 retry: retries.Retry | object | None = gapic_v1.method.DEFAULT,
52 timeout: float | None = None,
53 ) -> BatchWriteResponse:
54 """Writes the changes accumulated in this batch.
55
56 Write operations are not guaranteed to be applied in order and must not
57 contain multiple writes to any given document. Preferred over `commit`
58 for performance reasons if these conditions are acceptable.
59
60 Args:
61 retry (google.api_core.retry.Retry): Designation of what errors, if any,
62 should be retried. Defaults to a system-specified policy.
63 timeout (float): The timeout for this request. Defaults to a
64 system-specified value.
65
66 Returns:
67 :class:`google.cloud.firestore_v1.write.BatchWriteResponse`:
68 Container holding the write results corresponding to the changes
69 committed, returned in the same order as the changes were applied to
70 this batch. An individual write result contains an ``update_time``
71 field.
72 """
73 request, kwargs = self._prep_commit(retry, timeout)
74
75 _api = self._client._firestore_api
76 save_response: BatchWriteResponse = _api.batch_write(
77 request=request,
78 metadata=self._client._rpc_metadata,
79 **kwargs,
80 )
81
82 self._write_pbs = []
83 self.write_results = list(save_response.write_results)
84
85 return save_response
86
87 def _prep_commit(self, retry: retries.Retry | object | None, timeout: float | None):
88 request = {
89 "database": self._client._database_string,
90 "writes": self._write_pbs,
91 "labels": None,
92 }
93 kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
94 return request, kwargs