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