Skip to content

Commit

Permalink
Firestore: Implement firestore merge as a boolean
Browse files Browse the repository at this point in the history
  • Loading branch information
chemelnucfin committed Apr 6, 2018
1 parent 9561e4d commit fa1e5d0
Show file tree
Hide file tree
Showing 12 changed files with 58 additions and 185 deletions.
2 changes: 0 additions & 2 deletions firestore/google/cloud/firestore.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@
from google.cloud.firestore_v1beta1 import ExistsOption
from google.cloud.firestore_v1beta1 import GeoPoint
from google.cloud.firestore_v1beta1 import LastUpdateOption
from google.cloud.firestore_v1beta1 import MergeOption
from google.cloud.firestore_v1beta1 import Query
from google.cloud.firestore_v1beta1 import ReadAfterWriteError
from google.cloud.firestore_v1beta1 import SERVER_TIMESTAMP
Expand All @@ -47,7 +46,6 @@
'ExistsOption',
'GeoPoint',
'LastUpdateOption',
'MergeOption',
'Query',
'ReadAfterWriteError',
'SERVER_TIMESTAMP',
Expand Down
2 changes: 0 additions & 2 deletions firestore/google/cloud/firestore_v1beta1/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@
from google.cloud.firestore_v1beta1.client import Client
from google.cloud.firestore_v1beta1.client import ExistsOption
from google.cloud.firestore_v1beta1.client import LastUpdateOption
from google.cloud.firestore_v1beta1.client import MergeOption
from google.cloud.firestore_v1beta1.client import WriteOption
from google.cloud.firestore_v1beta1.collection import CollectionReference
from google.cloud.firestore_v1beta1.constants import DELETE_FIELD
Expand All @@ -48,7 +47,6 @@
'ExistsOption',
'GeoPoint',
'LastUpdateOption',
'MergeOption',
'Query',
'ReadAfterWriteError',
'SERVER_TIMESTAMP',
Expand Down
20 changes: 13 additions & 7 deletions firestore/google/cloud/firestore_v1beta1/_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -517,18 +517,18 @@ def encode_dict(values_dict):
}


def extract_field_paths(actual_data):
"""Extract field paths from actual data
def extract_field_paths(document_data):
"""Extract field paths from document data
Args:
actual_data (dict): The dictionary of the actual set data.
document_data (dict): The dictionary of the actual set data.
Returns:
List[~.firestore_v1beta1._helpers.FieldPath]:
A list of `FieldPath` instances from the actual data.
"""
field_paths = []
for field_name, value in six.iteritems(actual_data):
for field_name, value in six.iteritems(document_data):
if isinstance(value, dict):
sub_field_paths = extract_field_paths(value)
for sub_path in sub_field_paths:
Expand Down Expand Up @@ -897,7 +897,7 @@ def get_transform_pb(document_path, transform_paths):
)


def pbs_for_set(document_path, document_data, option):
def pbs_for_set(document_path, document_data, merge=False, exists=None):
"""Make ``Write`` protobufs for ``set()`` methods.
Args:
Expand All @@ -920,9 +920,15 @@ def pbs_for_set(document_path, document_data, option):
fields=encode_dict(actual_data),
),
)
if option is not None:
if exists is not None:
update_pb.current_document.CopyFrom(
common_pb2.Precondition(exists=exists))

if merge:
field_paths = extract_field_paths(document_data)
option.modify_write(update_pb, field_paths=field_paths)
field_paths = canonicalize_field_paths(field_paths)
mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
update_pb.update_mask.CopyFrom(mask)

write_pbs = [update_pb]
if transform_paths:
Expand Down
23 changes: 12 additions & 11 deletions firestore/google/cloud/firestore_v1beta1/batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,10 +57,11 @@ def create(self, reference, document_data):
document_data (dict): Property names and values to use for
creating a document.
"""
option = self._client.write_option(exists=False)
self.set(reference, document_data, option=option)
write_pbs = _helpers.pbs_for_set(
reference._document_path, document_data, merge=False, exists=False)
self._add_write_pbs(write_pbs)

def set(self, reference, document_data, option=None):
def set(self, reference, document_data, merge=False):
"""Add a "change" to replace a document.
See
Expand All @@ -69,16 +70,16 @@ def set(self, reference, document_data, option=None):
applied.
Args:
reference (~.firestore_v1beta1.document.DocumentReference): A
document reference that will have values set in this batch.
document_data (dict): Property names and values to use for
replacing a document.
option (Optional[~.firestore_v1beta1.client.WriteOption]): A
write option to make assertions / preconditions on the server
state of the document before applying changes.
reference (~.firestore_v1beta1.document.DocumentReference):
A document reference that will have values set in this batch.
document_data (dict):
Property names and values to use for replacing a document.
merge (Optional[bool]):
If True, apply merging instead of overwriting the state
of the document.
"""
write_pbs = _helpers.pbs_for_set(
reference._document_path, document_data, option)
reference._document_path, document_data, merge=merge)
self._add_write_pbs(write_pbs)

def update(self, reference, field_updates, option=None):
Expand Down
34 changes: 1 addition & 33 deletions firestore/google/cloud/firestore_v1beta1/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,14 +33,13 @@
from google.cloud.firestore_v1beta1.document import DocumentReference
from google.cloud.firestore_v1beta1.document import DocumentSnapshot
from google.cloud.firestore_v1beta1.gapic import firestore_client
from google.cloud.firestore_v1beta1.proto import common_pb2
from google.cloud.firestore_v1beta1.transaction import Transaction


DEFAULT_DATABASE = '(default)'
"""str: The default database used in a :class:`~.firestore.client.Client`."""
_BAD_OPTION_ERR = (
'Exactly one of ``last_update_time``, ``merge`` or ``exists`` '
'Exactly one of ``last_update_time``, ``exists`` '
'must be provided.')
_BAD_DOC_TEMPLATE = (
'Document {!r} appeared in response but was not present among references')
Expand Down Expand Up @@ -261,9 +260,6 @@ def write_option(**kwargs):
protobuf or directly.
* ``exists`` (:class:`bool`): Indicates if the document being modified
should already exist.
* ``merge`` (Any):
Indicates if the document should be merged.
**Note**: argument is ignored
Providing no argument would make the option have no effect (so
it is not allowed). Providing multiple would be an apparent
Expand All @@ -287,8 +283,6 @@ def write_option(**kwargs):
return LastUpdateOption(value)
elif name == 'exists':
return ExistsOption(value)
elif name == 'merge':
return MergeOption()
else:
extra = '{!r} was provided'.format(name)
raise TypeError(_BAD_OPTION_ERR, extra)
Expand Down Expand Up @@ -424,32 +418,6 @@ def modify_write(self, write_pb, **unused_kwargs):
write_pb.current_document.CopyFrom(current_doc)


class MergeOption(WriteOption):
"""Option used to merge on a write operation.
This will typically be created by
:meth:`~.firestore_v1beta1.client.Client.write_option`.
"""
def modify_write(
self, write_pb, field_paths=None, **unused_kwargs):
"""Modify a ``Write`` protobuf based on the state of this write option.
Args:
write_pb (google.cloud.firestore_v1beta1.types.Write): A
``Write`` protobuf instance to be modified with a precondition
determined by the state of this option.
field_paths (Sequence[str]):
The actual field names to use for replacing a document.
unused_kwargs (Dict[str, Any]): Keyword arguments accepted by
other subclasses that are unused here.
"""
if field_paths is None:
field_paths = []
field_paths = _helpers.canonicalize_field_paths(field_paths)
mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
write_pb.update_mask.CopyFrom(mask)


class ExistsOption(WriteOption):
"""Option used to assert existence on a write operation.
Expand Down
4 changes: 2 additions & 2 deletions firestore/google/cloud/firestore_v1beta1/document.py
Original file line number Diff line number Diff line change
Expand Up @@ -192,7 +192,7 @@ def create(self, document_data):
write_results = batch.commit()
return _first_write_result(write_results)

def set(self, document_data, option=None):
def set(self, document_data, merge=False):
"""Replace the current document in the Firestore database.
A write ``option`` can be specified to indicate preconditions of
Expand All @@ -219,7 +219,7 @@ def set(self, document_data, option=None):
result contains an ``update_time`` field.
"""
batch = self._client.batch()
batch.set(self, document_data, option=option)
batch.set(self, document_data, merge=merge)
write_results = batch.commit()
return _first_write_result(write_results)

Expand Down
52 changes: 10 additions & 42 deletions firestore/tests/system.py
Original file line number Diff line number Diff line change
Expand Up @@ -143,10 +143,9 @@ def test_document_set(client, cleanup):
snapshot = document.get()
assert snapshot.to_dict() is None

# 1. Use ``set()`` to create the document (using an option).
# 1. Use ``create()`` to create the document (using an option).
data1 = {'foo': 88}
option1 = client.write_option(exists=False)
write_result1 = document.set(data1, option=option1)
write_result1 = document.create(data1)
snapshot1 = document.get()
assert snapshot1.to_dict() == data1
# Make sure the update is what created the document.
Expand All @@ -162,30 +161,6 @@ def test_document_set(client, cleanup):
assert snapshot2.create_time == snapshot1.create_time
assert snapshot2.update_time == write_result2.update_time

# 3. Call ``set()`` with a valid "last timestamp" option.
data3 = {'skates': 88}
option3 = client.write_option(last_update_time=snapshot2.update_time)
write_result3 = document.set(data3, option=option3)
snapshot3 = document.get()
assert snapshot3.to_dict() == data3
# Make sure the create time hasn't changed.
assert snapshot3.create_time == snapshot1.create_time
assert snapshot3.update_time == write_result3.update_time

# 4. Call ``set()`` with invalid (in the past) "last timestamp" option.
assert_timestamp_less(option3._last_update_time, snapshot3.update_time)
with pytest.raises(FailedPrecondition):
document.set({'bad': 'time-past'}, option=option3)

# 5. Call ``set()`` with invalid (in the future) "last timestamp" option.
timestamp_pb = timestamp_pb2.Timestamp(
seconds=snapshot3.update_time.nanos + 120,
nanos=snapshot3.update_time.nanos,
)
option5 = client.write_option(last_update_time=timestamp_pb)
with pytest.raises(FailedPrecondition):
document.set({'bad': 'time-future'}, option=option5)


def test_document_integer_field(client, cleanup):
document_id = 'for-set' + unique_resource_id('-')
Expand All @@ -201,8 +176,7 @@ def test_document_integer_field(client, cleanup):
'7g': '8h',
'cd': '0j'}
}
option1 = client.write_option(exists=False)
document.set(data1, option=option1)
document.create(data1)

data2 = {'1a.ab': '4d', '6f.7g': '9h'}
option2 = client.write_option(exists=True)
Expand All @@ -225,30 +199,24 @@ def test_document_set_merge(client, cleanup):
# Add to clean-up before API request (in case ``set()`` fails).
cleanup(document)

# 0. Make sure the document doesn't exist yet using an option.
option0 = client.write_option(exists=True)
with pytest.raises(NotFound) as exc_info:
document.set({'no': 'way'}, option=option0)

assert exc_info.value.message.startswith(MISSING_DOCUMENT)
assert document_id in exc_info.value.message
# 0. Make sure the document doesn't exist yet
snapshot = document.get()
assert not snapshot.exists

# 1. Use ``set()`` to create the document (using an option).
data1 = {'name': 'Sam',
'address': {'city': 'SF',
'state': 'CA'}}
option1 = client.write_option(exists=False)
write_result1 = document.set(data1, option=option1)
write_result1 = document.create(data1)
snapshot1 = document.get()
assert snapshot1.to_dict() == data1
# Make sure the update is what created the document.
assert snapshot1.create_time == snapshot1.update_time
assert snapshot1.update_time == write_result1.update_time

# 2. Call ``set()`` again to overwrite (no option).
# 2. Call ``set()`` to merge
data2 = {'address': {'city': 'LA'}}
option2 = client.write_option(merge=True)
write_result2 = document.set(data2, option=option2)
write_result2 = document.set(data2, merge=True)
snapshot2 = document.get()
assert snapshot2.to_dict() == {'name': 'Sam',
'address': {'city': 'LA',
Expand Down Expand Up @@ -333,7 +301,7 @@ def test_update_document(client, cleanup):
)
option6 = client.write_option(last_update_time=timestamp_pb)
with pytest.raises(FailedPrecondition) as exc_info:
document.set({'bad': 'time-future'}, option=option6)
document.update({'bad': 'time-future'}, option=option6)


def check_snapshot(snapshot, document, data, write_result):
Expand Down
20 changes: 4 additions & 16 deletions firestore/tests/unit/test__helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -1397,8 +1397,7 @@ def _call_fut(document_path, document_data, option):

return pbs_for_set(document_path, document_data, option)

def _helper(self, option=None, do_transform=False, **write_kwargs):
from google.cloud.firestore_v1beta1.client import MergeOption
def _helper(self, merge=False, do_transform=False, **write_kwargs):
from google.cloud.firestore_v1beta1.constants import SERVER_TIMESTAMP
from google.cloud.firestore_v1beta1.gapic import enums
from google.cloud.firestore_v1beta1.proto import common_pb2
Expand All @@ -1420,7 +1419,7 @@ def _helper(self, option=None, do_transform=False, **write_kwargs):
if do_transform:
document_data[field_name3] = SERVER_TIMESTAMP

write_pbs = self._call_fut(document_path, document_data, option)
write_pbs = self._call_fut(document_path, document_data, merge)

expected_update_pb = write_pb2.Write(
update=document_pb2.Document(
Expand All @@ -1434,7 +1433,7 @@ def _helper(self, option=None, do_transform=False, **write_kwargs):
)
expected_pbs = [expected_update_pb]

if isinstance(option, MergeOption):
if merge:
field_paths = [field_name1, field_name2]
mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
expected_pbs[0].update_mask.CopyFrom(mask)
Expand All @@ -1459,19 +1458,8 @@ def _helper(self, option=None, do_transform=False, **write_kwargs):
def test_without_option(self):
self._helper()

def test_with_exists_option(self):
from google.cloud.firestore_v1beta1.proto import common_pb2
from google.cloud.firestore_v1beta1.client import ExistsOption

option = ExistsOption(True)
precondition = common_pb2.Precondition(exists=True)
self._helper(option=option, current_document=precondition)

def test_with_merge_option(self):
from google.cloud.firestore_v1beta1.client import MergeOption

option = MergeOption()
self._helper(option=option)
self._helper(merge=True)

def test_update_and_transform(self):
self._helper(do_transform=True)
Expand Down
4 changes: 1 addition & 3 deletions firestore/tests/unit/test_batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,6 @@ def test_set(self):
def test_set_merge(self):
from google.cloud.firestore_v1beta1.proto import document_pb2
from google.cloud.firestore_v1beta1.proto import write_pb2
from google.cloud.firestore_v1beta1.client import MergeOption

client = _make_client()
batch = self._make_one(client)
Expand All @@ -103,8 +102,7 @@ def test_set_merge(self):
field = 'zapzap'
value = u'meadows and flowers'
document_data = {field: value}
option = MergeOption()
ret_val = batch.set(reference, document_data, option)
ret_val = batch.set(reference, document_data, merge=True)
self.assertIsNone(ret_val)
new_write_pb = write_pb2.Write(
update=document_pb2.Document(
Expand Down
Loading

0 comments on commit fa1e5d0

Please sign in to comment.