Skip to content

Commit 8b2dfd3

Browse files
authored
refactor: create base batch class (#67)
towards #65 this PR is staged on top of #66 as later refactor changes will be cross-class; I encourage reviewing in sequential order
1 parent 02840bf commit 8b2dfd3

File tree

4 files changed

+307
-224
lines changed

4 files changed

+307
-224
lines changed
Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
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+
17+
18+
from google.cloud.firestore_v1 import _helpers
19+
20+
21+
class BaseWriteBatch(object):
22+
"""Accumulate write operations to be sent in a batch.
23+
24+
This has the same set of methods for write operations that
25+
:class:`~google.cloud.firestore_v1.document.DocumentReference` does,
26+
e.g. :meth:`~google.cloud.firestore_v1.document.DocumentReference.create`.
27+
28+
Args:
29+
client (:class:`~google.cloud.firestore_v1.client.Client`):
30+
The client that created this batch.
31+
"""
32+
33+
def __init__(self, client):
34+
self._client = client
35+
self._write_pbs = []
36+
self.write_results = None
37+
self.commit_time = None
38+
39+
def _add_write_pbs(self, write_pbs):
40+
"""Add `Write`` protobufs to this transaction.
41+
42+
This method intended to be over-ridden by subclasses.
43+
44+
Args:
45+
write_pbs (List[google.cloud.proto.firestore.v1.\
46+
write_pb2.Write]): A list of write protobufs to be added.
47+
"""
48+
self._write_pbs.extend(write_pbs)
49+
50+
def create(self, reference, document_data):
51+
"""Add a "change" to this batch to create a document.
52+
53+
If the document given by ``reference`` already exists, then this
54+
batch will fail when :meth:`commit`-ed.
55+
56+
Args:
57+
reference (:class:`~google.cloud.firestore_v1.document.DocumentReference`):
58+
A document reference to be created in this batch.
59+
document_data (dict): Property names and values to use for
60+
creating a document.
61+
"""
62+
write_pbs = _helpers.pbs_for_create(reference._document_path, document_data)
63+
self._add_write_pbs(write_pbs)
64+
65+
def set(self, reference, document_data, merge=False):
66+
"""Add a "change" to replace a document.
67+
68+
See
69+
:meth:`google.cloud.firestore_v1.document.DocumentReference.set` for
70+
more information on how ``option`` determines how the change is
71+
applied.
72+
73+
Args:
74+
reference (:class:`~google.cloud.firestore_v1.document.DocumentReference`):
75+
A document reference that will have values set in this batch.
76+
document_data (dict):
77+
Property names and values to use for replacing a document.
78+
merge (Optional[bool] or Optional[List<apispec>]):
79+
If True, apply merging instead of overwriting the state
80+
of the document.
81+
"""
82+
if merge is not False:
83+
write_pbs = _helpers.pbs_for_set_with_merge(
84+
reference._document_path, document_data, merge
85+
)
86+
else:
87+
write_pbs = _helpers.pbs_for_set_no_merge(
88+
reference._document_path, document_data
89+
)
90+
91+
self._add_write_pbs(write_pbs)
92+
93+
def update(self, reference, field_updates, option=None):
94+
"""Add a "change" to update a document.
95+
96+
See
97+
:meth:`google.cloud.firestore_v1.document.DocumentReference.update`
98+
for more information on ``field_updates`` and ``option``.
99+
100+
Args:
101+
reference (:class:`~google.cloud.firestore_v1.document.DocumentReference`):
102+
A document reference that will be updated in this batch.
103+
field_updates (dict):
104+
Field names or paths to update and values to update with.
105+
option (Optional[:class:`~google.cloud.firestore_v1.client.WriteOption`]):
106+
A write option to make assertions / preconditions on the server
107+
state of the document before applying changes.
108+
"""
109+
if option.__class__.__name__ == "ExistsOption":
110+
raise ValueError("you must not pass an explicit write option to " "update.")
111+
write_pbs = _helpers.pbs_for_update(
112+
reference._document_path, field_updates, option
113+
)
114+
self._add_write_pbs(write_pbs)
115+
116+
def delete(self, reference, option=None):
117+
"""Add a "change" to delete a document.
118+
119+
See
120+
:meth:`google.cloud.firestore_v1.document.DocumentReference.delete`
121+
for more information on how ``option`` determines how the change is
122+
applied.
123+
124+
Args:
125+
reference (:class:`~google.cloud.firestore_v1.document.DocumentReference`):
126+
A document reference that will be deleted in this batch.
127+
option (Optional[:class:`~google.cloud.firestore_v1.client.WriteOption`]):
128+
A write option to make assertions / preconditions on the server
129+
state of the document before applying changes.
130+
"""
131+
write_pb = _helpers.pb_for_delete(reference._document_path, option)
132+
self._add_write_pbs([write_pb])

packages/google-cloud-firestore/google/cloud/firestore_v1/batch.py

Lines changed: 3 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,10 @@
1515
"""Helpers for batch requests to the Google Cloud Firestore API."""
1616

1717

18-
from google.cloud.firestore_v1 import _helpers
18+
from google.cloud.firestore_v1.base_batch import BaseWriteBatch
1919

2020

21-
class WriteBatch(object):
21+
class WriteBatch(BaseWriteBatch):
2222
"""Accumulate write operations to be sent in a batch.
2323
2424
This has the same set of methods for write operations that
@@ -31,105 +31,7 @@ class WriteBatch(object):
3131
"""
3232

3333
def __init__(self, client):
34-
self._client = client
35-
self._write_pbs = []
36-
self.write_results = None
37-
self.commit_time = None
38-
39-
def _add_write_pbs(self, write_pbs):
40-
"""Add `Write`` protobufs to this transaction.
41-
42-
This method intended to be over-ridden by subclasses.
43-
44-
Args:
45-
write_pbs (List[google.cloud.proto.firestore.v1.\
46-
write_pb2.Write]): A list of write protobufs to be added.
47-
"""
48-
self._write_pbs.extend(write_pbs)
49-
50-
def create(self, reference, document_data):
51-
"""Add a "change" to this batch to create a document.
52-
53-
If the document given by ``reference`` already exists, then this
54-
batch will fail when :meth:`commit`-ed.
55-
56-
Args:
57-
reference (:class:`~google.cloud.firestore_v1.document.DocumentReference`):
58-
A document reference to be created in this batch.
59-
document_data (dict): Property names and values to use for
60-
creating a document.
61-
"""
62-
write_pbs = _helpers.pbs_for_create(reference._document_path, document_data)
63-
self._add_write_pbs(write_pbs)
64-
65-
def set(self, reference, document_data, merge=False):
66-
"""Add a "change" to replace a document.
67-
68-
See
69-
:meth:`google.cloud.firestore_v1.document.DocumentReference.set` for
70-
more information on how ``option`` determines how the change is
71-
applied.
72-
73-
Args:
74-
reference (:class:`~google.cloud.firestore_v1.document.DocumentReference`):
75-
A document reference that will have values set in this batch.
76-
document_data (dict):
77-
Property names and values to use for replacing a document.
78-
merge (Optional[bool] or Optional[List<apispec>]):
79-
If True, apply merging instead of overwriting the state
80-
of the document.
81-
"""
82-
if merge is not False:
83-
write_pbs = _helpers.pbs_for_set_with_merge(
84-
reference._document_path, document_data, merge
85-
)
86-
else:
87-
write_pbs = _helpers.pbs_for_set_no_merge(
88-
reference._document_path, document_data
89-
)
90-
91-
self._add_write_pbs(write_pbs)
92-
93-
def update(self, reference, field_updates, option=None):
94-
"""Add a "change" to update a document.
95-
96-
See
97-
:meth:`google.cloud.firestore_v1.document.DocumentReference.update`
98-
for more information on ``field_updates`` and ``option``.
99-
100-
Args:
101-
reference (:class:`~google.cloud.firestore_v1.document.DocumentReference`):
102-
A document reference that will be updated in this batch.
103-
field_updates (dict):
104-
Field names or paths to update and values to update with.
105-
option (Optional[:class:`~google.cloud.firestore_v1.client.WriteOption`]):
106-
A write option to make assertions / preconditions on the server
107-
state of the document before applying changes.
108-
"""
109-
if option.__class__.__name__ == "ExistsOption":
110-
raise ValueError("you must not pass an explicit write option to " "update.")
111-
write_pbs = _helpers.pbs_for_update(
112-
reference._document_path, field_updates, option
113-
)
114-
self._add_write_pbs(write_pbs)
115-
116-
def delete(self, reference, option=None):
117-
"""Add a "change" to delete a document.
118-
119-
See
120-
:meth:`google.cloud.firestore_v1.document.DocumentReference.delete`
121-
for more information on how ``option`` determines how the change is
122-
applied.
123-
124-
Args:
125-
reference (:class:`~google.cloud.firestore_v1.document.DocumentReference`):
126-
A document reference that will be deleted in this batch.
127-
option (Optional[:class:`~google.cloud.firestore_v1.client.WriteOption`]):
128-
A write option to make assertions / preconditions on the server
129-
state of the document before applying changes.
130-
"""
131-
write_pb = _helpers.pb_for_delete(reference._document_path, option)
132-
self._add_write_pbs([write_pb])
34+
super(WriteBatch, self).__init__(client=client)
13335

13436
def commit(self):
13537
"""Commit the changes accumulated in this batch.

0 commit comments

Comments
 (0)