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

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

16 

17import abc 

18from typing import Dict, Union 

19 

20# Types needed only for Type Hints 

21from google.api_core import retry as retries 

22from google.cloud.firestore_v1 import _helpers 

23from google.cloud.firestore_v1.base_document import BaseDocumentReference 

24 

25 

26class BaseBatch(metaclass=abc.ABCMeta): 

27 """Accumulate write operations to be sent in a batch. 

28 

29 This has the same set of methods for write operations that 

30 :class:`~google.cloud.firestore_v1.document.DocumentReference` does, 

31 e.g. :meth:`~google.cloud.firestore_v1.document.DocumentReference.create`. 

32 

33 Args: 

34 client (:class:`~google.cloud.firestore_v1.client.Client`): 

35 The client that created this batch. 

36 """ 

37 

38 def __init__(self, client) -> None: 

39 self._client = client 

40 self._write_pbs = [] 

41 self._document_references: Dict[str, BaseDocumentReference] = {} 

42 self.write_results = None 

43 self.commit_time = None 

44 

45 def __len__(self): 

46 return len(self._document_references) 

47 

48 def __contains__(self, reference: BaseDocumentReference): 

49 return reference._document_path in self._document_references 

50 

51 def _add_write_pbs(self, write_pbs: list) -> None: 

52 """Add `Write`` protobufs to this transaction. 

53 

54 This method intended to be over-ridden by subclasses. 

55 

56 Args: 

57 write_pbs (List[google.cloud.proto.firestore.v1.\ 

58 write_pb2.Write]): A list of write protobufs to be added. 

59 """ 

60 self._write_pbs.extend(write_pbs) 

61 

62 @abc.abstractmethod 

63 def commit(self): 

64 """Sends all accumulated write operations to the server. The details of this 

65 write depend on the implementing class.""" 

66 raise NotImplementedError() 

67 

68 def create(self, reference: BaseDocumentReference, document_data: dict) -> None: 

69 """Add a "change" to this batch to create a document. 

70 

71 If the document given by ``reference`` already exists, then this 

72 batch will fail when :meth:`commit`-ed. 

73 

74 Args: 

75 reference (:class:`~google.cloud.firestore_v1.document.DocumentReference`): 

76 A document reference to be created in this batch. 

77 document_data (dict): Property names and values to use for 

78 creating a document. 

79 """ 

80 write_pbs = _helpers.pbs_for_create(reference._document_path, document_data) 

81 self._document_references[reference._document_path] = reference 

82 self._add_write_pbs(write_pbs) 

83 

84 def set( 

85 self, 

86 reference: BaseDocumentReference, 

87 document_data: dict, 

88 merge: Union[bool, list] = False, 

89 ) -> None: 

90 """Add a "change" to replace a document. 

91 

92 See 

93 :meth:`google.cloud.firestore_v1.document.DocumentReference.set` for 

94 more information on how ``option`` determines how the change is 

95 applied. 

96 

97 Args: 

98 reference (:class:`~google.cloud.firestore_v1.document.DocumentReference`): 

99 A document reference that will have values set in this batch. 

100 document_data (dict): 

101 Property names and values to use for replacing a document. 

102 merge (Optional[bool] or Optional[List<apispec>]): 

103 If True, apply merging instead of overwriting the state 

104 of the document. 

105 """ 

106 if merge is not False: 

107 write_pbs = _helpers.pbs_for_set_with_merge( 

108 reference._document_path, document_data, merge 

109 ) 

110 else: 

111 write_pbs = _helpers.pbs_for_set_no_merge( 

112 reference._document_path, document_data 

113 ) 

114 

115 self._document_references[reference._document_path] = reference 

116 self._add_write_pbs(write_pbs) 

117 

118 def update( 

119 self, 

120 reference: BaseDocumentReference, 

121 field_updates: dict, 

122 option: _helpers.WriteOption = None, 

123 ) -> None: 

124 """Add a "change" to update a document. 

125 

126 See 

127 :meth:`google.cloud.firestore_v1.document.DocumentReference.update` 

128 for more information on ``field_updates`` and ``option``. 

129 

130 Args: 

131 reference (:class:`~google.cloud.firestore_v1.document.DocumentReference`): 

132 A document reference that will be updated in this batch. 

133 field_updates (dict): 

134 Field names or paths to update and values to update with. 

135 option (Optional[:class:`~google.cloud.firestore_v1.client.WriteOption`]): 

136 A write option to make assertions / preconditions on the server 

137 state of the document before applying changes. 

138 """ 

139 if option.__class__.__name__ == "ExistsOption": 

140 raise ValueError("you must not pass an explicit write option to " "update.") 

141 write_pbs = _helpers.pbs_for_update( 

142 reference._document_path, field_updates, option 

143 ) 

144 self._document_references[reference._document_path] = reference 

145 self._add_write_pbs(write_pbs) 

146 

147 def delete( 

148 self, reference: BaseDocumentReference, option: _helpers.WriteOption = None 

149 ) -> None: 

150 """Add a "change" to delete a document. 

151 

152 See 

153 :meth:`google.cloud.firestore_v1.document.DocumentReference.delete` 

154 for more information on how ``option`` determines how the change is 

155 applied. 

156 

157 Args: 

158 reference (:class:`~google.cloud.firestore_v1.document.DocumentReference`): 

159 A document reference that will be deleted in this batch. 

160 option (Optional[:class:`~google.cloud.firestore_v1.client.WriteOption`]): 

161 A write option to make assertions / preconditions on the server 

162 state of the document before applying changes. 

163 """ 

164 write_pb = _helpers.pb_for_delete(reference._document_path, option) 

165 self._document_references[reference._document_path] = reference 

166 self._add_write_pbs([write_pb]) 

167 

168 

169class BaseWriteBatch(BaseBatch): 

170 """Base class for a/sync implementations of the `commit` RPC. `commit` is useful 

171 for lower volumes or when the order of write operations is important.""" 

172 

173 def _prep_commit(self, retry: retries.Retry, timeout: float): 

174 """Shared setup for async/sync :meth:`commit`.""" 

175 request = { 

176 "database": self._client._database_string, 

177 "writes": self._write_pbs, 

178 "transaction": None, 

179 } 

180 kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) 

181 return request, kwargs