Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/google/cloud/firestore_v1/batch.py: 47%

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.""" 

16 

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, retry: retries.Retry = gapic_v1.method.DEFAULT, timeout: float = None 

42 ) -> list: 

43 """Commit the changes accumulated in this batch. 

44 

45 Args: 

46 retry (google.api_core.retry.Retry): Designation of what errors, if any, 

47 should be retried. Defaults to a system-specified policy. 

48 timeout (float): The timeout for this request. Defaults to a 

49 system-specified value. 

50 

51 Returns: 

52 List[:class:`google.cloud.proto.firestore.v1.write.WriteResult`, ...]: 

53 The write results corresponding to the changes committed, returned 

54 in the same order as the changes were applied to this batch. A 

55 write result contains an ``update_time`` field. 

56 """ 

57 request, kwargs = self._prep_commit(retry, timeout) 

58 

59 commit_response = self._client._firestore_api.commit( 

60 request=request, 

61 metadata=self._client._rpc_metadata, 

62 **kwargs, 

63 ) 

64 

65 self._write_pbs = [] 

66 self.write_results = results = list(commit_response.write_results) 

67 self.commit_time = commit_response.commit_time 

68 

69 return results 

70 

71 def __enter__(self): 

72 return self 

73 

74 def __exit__(self, exc_type, exc_value, traceback): 

75 if exc_type is None: 

76 self.commit()