Skip to content

Commit 77795e1

Browse files
committed
Firestore: MergeOption
1 parent 9d7b462 commit 77795e1

File tree

5 files changed

+123
-14
lines changed

5 files changed

+123
-14
lines changed

firestore/google/cloud/firestore_v1beta1/_helpers.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,9 @@
2828
import grpc
2929
import six
3030

31+
from google.cloud import exceptions
3132
from google.cloud._helpers import _datetime_to_pb_timestamp
3233
from google.cloud._helpers import _pb_timestamp_to_datetime
33-
from google.cloud import exceptions
34-
3534
from google.cloud.firestore_v1beta1 import constants
3635
from google.cloud.firestore_v1beta1.gapic import enums
3736
from google.cloud.firestore_v1beta1.proto import common_pb2
@@ -813,15 +812,14 @@ def pbs_for_set(document_path, document_data, option):
813812
or two ``Write`` protobuf instances for ``set()``.
814813
"""
815814
transform_paths, actual_data = remove_server_timestamp(document_data)
816-
817815
update_pb = write_pb2.Write(
818816
update=document_pb2.Document(
819817
name=document_path,
820818
fields=encode_dict(actual_data),
821819
),
822820
)
823821
if option is not None:
824-
option.modify_write(update_pb)
822+
option.modify_write(update_pb, actual_data=actual_data, path=document_path)
825823

826824
write_pbs = [update_pb]
827825
if transform_paths:
@@ -866,7 +864,7 @@ def pbs_for_update(client, document_path, field_updates, option):
866864
update_mask=common_pb2.DocumentMask(field_paths=sorted(field_paths)),
867865
)
868866
# Due to the default, we don't have to check if ``None``.
869-
option.modify_write(update_pb)
867+
option.modify_write(update_pb, field_paths=field_paths)
870868
write_pbs = [update_pb]
871869

872870
if transform_paths:

firestore/google/cloud/firestore_v1beta1/client.py

Lines changed: 47 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,9 @@
3737
from google.cloud.firestore_v1beta1.document import DocumentSnapshot
3838
from google.cloud.firestore_v1beta1.gapic import firestore_client
3939
from google.cloud.firestore_v1beta1.transaction import Transaction
40+
from google.cloud.firestore_v1beta1.proto import common_pb2
41+
from google.cloud.firestore_v1beta1.proto import document_pb2
42+
from google.cloud.firestore_v1beta1.proto import write_pb2
4043

4144

4245
DEFAULT_DATABASE = '(default)'
@@ -251,15 +254,18 @@ def write_option(**kwargs):
251254
:meth:`~.DocumentReference.update` and
252255
:meth:`~.DocumentReference.delete`.
253256
254-
One of the following two keyword arguments must be provided:
257+
One of the following keyword arguments must be provided:
255258
256259
* ``last_update_time`` (:class:`google.protobuf.timestamp_pb2.\
257-
Timestamp`): A timestamp. When set, the target document must exist
258-
and have been last updated at that time. Protobuf ``update_time``
259-
timestamps are typically returned from methods that perform write
260-
operations as part of a "write result" protobuf or directly.
260+
Timestamp`): A timestamp. When set, the target document must exist
261+
and have been last updated at that time. Protobuf ``update_time``
262+
timestamps are typically returned from methods that perform write
263+
operations as part of a "write result" protobuf or directly.
261264
* ``exists`` (:class:`bool`): Indicates if the document being modified
262-
should already exist.
265+
should already exist.
266+
* ``merge`` (Any):
267+
Indicates if the document should be merged.
268+
**Note**: argument is ignored
263269
264270
Providing no argument would make the option have no effect (so
265271
it is not allowed). Providing multiple would be an apparent
@@ -283,6 +289,8 @@ def write_option(**kwargs):
283289
return LastUpdateOption(value)
284290
elif name == 'exists':
285291
return ExistsOption(value)
292+
elif name == 'merge':
293+
return MergeOption()
286294
else:
287295
extra = '{!r} was provided'.format(name)
288296
raise TypeError(_BAD_OPTION_ERR, extra)
@@ -418,6 +426,39 @@ def modify_write(self, write_pb, **unused_kwargs):
418426
write_pb.current_document.CopyFrom(current_doc)
419427

420428

429+
class MergeOption(WriteOption):
430+
"""Option used to merge on a write operation.
431+
432+
This will typically be created by
433+
:meth:`~.firestore_v1beta1.client.Client.write_option`.
434+
"""
435+
def modify_write(self, write_pb, actual_data=None, path=None, **unused_kwargs):
436+
"""Modify a ``Write`` protobuf based on the state of this write option.
437+
438+
Args:
439+
write_pb (google.cloud.firestore_v1beta1.types.Write): A
440+
``Write`` protobuf instance to be modified with a precondition
441+
determined by the state of this option.
442+
actual_data (dict):
443+
The actual field names and values to use for replacing a
444+
document.
445+
path (str): A fully-qualified document_path
446+
unused_kwargs (Dict[str, Any]): Keyword arguments accepted by
447+
other subclasses that are unused here.
448+
"""
449+
actual_data, field_paths = _helpers.FieldPathHelper.to_field_paths(actual_data)
450+
doc = document_pb2.Document(
451+
name=path,
452+
fields=_helpers.encode_dict(actual_data)
453+
)
454+
write = write_pb2.Write(
455+
update=doc,
456+
)
457+
write_pb.CopyFrom(write)
458+
mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
459+
write_pb.update_mask.CopyFrom(mask)
460+
461+
421462
class ExistsOption(WriteOption):
422463
"""Option used to assert existence on a write operation.
423464

firestore/tests/system.py

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,47 @@ def test_document_set(client, cleanup):
194194
assert exc_to_code(exc_info.value.cause) == StatusCode.FAILED_PRECONDITION
195195

196196

197+
def test_document_set_merge(client, cleanup):
198+
document_id = 'for-set' + unique_resource_id('-')
199+
document = client.document('i-did-it', document_id)
200+
# Add to clean-up before API request (in case ``set()`` fails).
201+
cleanup(document)
202+
203+
# 0. Make sure the document doesn't exist yet using an option.
204+
option0 = client.write_option(exists=True)
205+
with pytest.raises(NotFound) as exc_info:
206+
document.set({'no': 'way'}, option=option0)
207+
208+
assert exc_info.value.message.startswith(MISSING_DOCUMENT)
209+
assert document_id in exc_info.value.message
210+
211+
# 1. Use ``set()`` to create the document (using an option).
212+
data1 = {'name': 'Sam',
213+
'address': {'city': 'SF',
214+
'state': 'CA'}
215+
}
216+
option1 = client.write_option(exists=False)
217+
write_result1 = document.set(data1, option=option1)
218+
snapshot1 = document.get()
219+
assert snapshot1.to_dict() == data1
220+
# Make sure the update is what created the document.
221+
assert snapshot1.create_time == snapshot1.update_time
222+
assert snapshot1.update_time == write_result1.update_time
223+
224+
# 2. Call ``set()`` again to overwrite (no option).
225+
data2 = {'address.city': 'LA'}
226+
option2 = client.write_option(merge=True)
227+
write_result2 = document.set(data2, option=option2)
228+
snapshot2 = document.get()
229+
assert snapshot2.to_dict() == {'name': 'Sam',
230+
'address': {'city': 'LA',
231+
'state': 'CA'}
232+
}
233+
# Make sure the create time hasn't changed.
234+
assert snapshot2.create_time == snapshot1.create_time
235+
assert snapshot2.update_time == write_result2.update_time
236+
237+
197238
def test_update_document(client, cleanup):
198239
document_id = 'for-update' + unique_resource_id('-')
199240
document = client.document('made', document_id)
@@ -207,7 +248,6 @@ def test_update_document(client, cleanup):
207248
assert document_id in exc_info.value.message
208249

209250
# 1. Try to update before the document exists (now with an option).
210-
211251
option1 = client.write_option(exists=True)
212252
with pytest.raises(NotFound) as exc_info:
213253
document.update({'still': 'not-there'}, option=option1)

firestore/tests/unit/test_batch.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,33 @@ def test_set(self):
9393
)
9494
self.assertEqual(batch._write_pbs, [new_write_pb])
9595

96+
def test_set_merge(self):
97+
from google.cloud.firestore_v1beta1.proto import document_pb2
98+
from google.cloud.firestore_v1beta1.proto import write_pb2
99+
from google.cloud.firestore_v1beta1.client import MergeOption
100+
101+
client = _make_client()
102+
batch = self._make_one(client)
103+
self.assertEqual(batch._write_pbs, [])
104+
105+
reference = client.document('another', 'one')
106+
field = 'zapzap'
107+
value = u'meadows and flowers'
108+
document_data = {field: value}
109+
option = MergeOption()
110+
ret_val = batch.set(reference, document_data, option)
111+
self.assertIsNone(ret_val)
112+
new_write_pb = write_pb2.Write(
113+
update=document_pb2.Document(
114+
name=reference._document_path,
115+
fields={
116+
field: _value_pb(string_value=value),
117+
},
118+
),
119+
update_mask={'field_paths':[field]}
120+
)
121+
self.assertEqual(batch._write_pbs, [new_write_pb])
122+
96123
def test_update(self):
97124
from google.cloud.firestore_v1beta1.proto import common_pb2
98125
from google.cloud.firestore_v1beta1.proto import document_pb2

firestore/tests/unit/test_document.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -249,6 +249,7 @@ def _write_pb_for_set(document_path, document_data):
249249
)
250250

251251
def _set_helper(self, **option_kwargs):
252+
from google.cloud.firestore_v1beta1._helpers import FieldPathHelper
252253
# Create a minimal fake GAPIC with a dummy response.
253254
firestore_api = mock.Mock(spec=['commit'])
254255
commit_response = mock.Mock(
@@ -277,16 +278,18 @@ def _set_helper(self, **option_kwargs):
277278
self.assertIs(write_result, mock.sentinel.write_result)
278279
write_pb = self._write_pb_for_set(
279280
document._document_path, document_data)
281+
update_values, field_paths = FieldPathHelper.to_field_paths(document_data)
282+
280283
if option is not None:
281-
option.modify_write(write_pb)
284+
option.modify_write(write_pb, field_paths=field_paths)
282285
firestore_api.commit.assert_called_once_with(
283286
client._database_string, [write_pb], transaction=None,
284287
options=client._call_options)
285288

286289
def test_set(self):
287290
self._set_helper()
288291

289-
def test_set_with_option(self):
292+
def test_set_exists(self):
290293
self._set_helper(exists=True)
291294

292295
@staticmethod

0 commit comments

Comments
 (0)