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"""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.base_batch import BaseWriteBatch
21
22
23class WriteBatch(BaseWriteBatch):
24 """Accumulate write operations to be sent in a batch. Use this over
25 `BulkWriteBatch` for lower volumes or when the order of operations
26 within a given batch is important.
27
28 This has the same set of methods for write operations that
29 :class:`~google.cloud.firestore_v1.document.DocumentReference` does,
30 e.g. :meth:`~google.cloud.firestore_v1.document.DocumentReference.create`.
31
32 Args:
33 client (:class:`~google.cloud.firestore_v1.client.Client`):
34 The client that created this batch.
35 """
36
37 def __init__(self, client) -> None:
38 super(WriteBatch, self).__init__(client=client)
39
40 def commit(
41 self,
42 retry: retries.Retry | object | None = gapic_v1.method.DEFAULT,
43 timeout: float | None = None,
44 ) -> list:
45 """Commit the changes accumulated in this batch.
46
47 Args:
48 retry (google.api_core.retry.Retry): Designation of what errors, if any,
49 should be retried. Defaults to a system-specified policy.
50 timeout (float): The timeout for this request. Defaults to a
51 system-specified value.
52
53 Returns:
54 List[:class:`google.cloud.firestore_v1.write.WriteResult`, ...]:
55 The write results corresponding to the changes committed, returned
56 in the same order as the changes were applied to this batch. A
57 write result contains an ``update_time`` field.
58 """
59 request, kwargs = self._prep_commit(retry, timeout)
60
61 commit_response = self._client._firestore_api.commit(
62 request=request,
63 metadata=self._client._rpc_metadata,
64 **kwargs,
65 )
66
67 self._write_pbs = []
68 self.write_results = results = list(commit_response.write_results)
69 self.commit_time = commit_response.commit_time
70
71 return results
72
73 def __enter__(self):
74 return self
75
76 def __exit__(self, exc_type, exc_value, traceback):
77 if exc_type is None:
78 self.commit()