Skip to content
This repository was archived by the owner on Mar 2, 2026. It is now read-only.

Commit bd51726

Browse files
committed
refactor: create base transaction class
1 parent b4a8eb9 commit bd51726

4 files changed

Lines changed: 317 additions & 167 deletions

File tree

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
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 applying Google Cloud Firestore changes in a transaction."""
16+
17+
18+
from google.cloud.firestore_v1 import types
19+
20+
MAX_ATTEMPTS = 5
21+
"""int: Default number of transaction attempts (with retries)."""
22+
_CANT_BEGIN = "The transaction has already begun. Current transaction ID: {!r}."
23+
_MISSING_ID_TEMPLATE = "The transaction has no transaction ID, so it cannot be {}."
24+
_CANT_ROLLBACK = _MISSING_ID_TEMPLATE.format("rolled back")
25+
_CANT_COMMIT = _MISSING_ID_TEMPLATE.format("committed")
26+
_WRITE_READ_ONLY = "Cannot perform write operation in read-only transaction."
27+
_INITIAL_SLEEP = 1.0
28+
"""float: Initial "max" for sleep interval. To be used in :func:`_sleep`."""
29+
_MAX_SLEEP = 30.0
30+
"""float: Eventual "max" sleep time. To be used in :func:`_sleep`."""
31+
_MULTIPLIER = 2.0
32+
"""float: Multiplier for exponential backoff. To be used in :func:`_sleep`."""
33+
_EXCEED_ATTEMPTS_TEMPLATE = "Failed to commit transaction in {:d} attempts."
34+
_CANT_RETRY_READ_ONLY = "Only read-write transactions can be retried."
35+
36+
37+
class BaseTransaction(object):
38+
"""Accumulate read-and-write operations to be sent in a transaction.
39+
40+
Args:
41+
max_attempts (Optional[int]): The maximum number of attempts for
42+
the transaction (i.e. allowing retries). Defaults to
43+
:attr:`~google.cloud.firestore_v1.transaction.MAX_ATTEMPTS`.
44+
read_only (Optional[bool]): Flag indicating if the transaction
45+
should be read-only or should allow writes. Defaults to
46+
:data:`False`.
47+
"""
48+
49+
def __init__(self, max_attempts=MAX_ATTEMPTS, read_only=False):
50+
self._max_attempts = max_attempts
51+
self._read_only = read_only
52+
self._id = None
53+
54+
def _add_write_pbs(self, write_pbs):
55+
raise NotImplementedError
56+
57+
def _options_protobuf(self, retry_id):
58+
"""Convert the current object to protobuf.
59+
60+
The ``retry_id`` value is used when retrying a transaction that
61+
failed (e.g. due to contention). It is intended to be the "first"
62+
transaction that failed (i.e. if multiple retries are needed).
63+
64+
Args:
65+
retry_id (Union[bytes, NoneType]): Transaction ID of a transaction
66+
to be retried.
67+
68+
Returns:
69+
Optional[google.cloud.firestore_v1.types.TransactionOptions]:
70+
The protobuf ``TransactionOptions`` if ``read_only==True`` or if
71+
there is a transaction ID to be retried, else :data:`None`.
72+
73+
Raises:
74+
ValueError: If ``retry_id`` is not :data:`None` but the
75+
transaction is read-only.
76+
"""
77+
if retry_id is not None:
78+
if self._read_only:
79+
raise ValueError(_CANT_RETRY_READ_ONLY)
80+
81+
return types.TransactionOptions(
82+
read_write=types.TransactionOptions.ReadWrite(
83+
retry_transaction=retry_id
84+
)
85+
)
86+
elif self._read_only:
87+
return types.TransactionOptions(
88+
read_only=types.TransactionOptions.ReadOnly()
89+
)
90+
else:
91+
return None
92+
93+
@property
94+
def in_progress(self):
95+
"""Determine if this transaction has already begun.
96+
97+
Returns:
98+
bool: Indicates if the transaction has started.
99+
"""
100+
return self._id is not None
101+
102+
@property
103+
def id(self):
104+
"""Get the current transaction ID.
105+
106+
Returns:
107+
Optional[bytes]: The transaction ID (or :data:`None` if the
108+
current transaction is not in progress).
109+
"""
110+
return self._id
111+
112+
def _clean_up(self):
113+
"""Clean up the instance after :meth:`_rollback`` or :meth:`_commit``.
114+
115+
This intended to occur on success or failure of the associated RPCs.
116+
"""
117+
self._write_pbs = []
118+
self._id = None
119+
120+
def _begin(self, retry_id=None):
121+
raise NotImplementedError
122+
123+
def _rollback(self):
124+
raise NotImplementedError
125+
126+
def _commit(self):
127+
raise NotImplementedError
128+
129+
def get_all(self, references):
130+
raise NotImplementedError
131+
132+
def get(self, ref_or_query):
133+
raise NotImplementedError
134+
135+
136+
class _BaseTransactional(object):
137+
"""Provide a callable object to use as a transactional decorater.
138+
139+
This is surfaced via
140+
:func:`~google.cloud.firestore_v1.transaction.transactional`.
141+
142+
Args:
143+
to_wrap (Callable[[:class:`~google.cloud.firestore_v1.transaction.Transaction`, ...], Any]):
144+
A callable that should be run (and retried) in a transaction.
145+
"""
146+
147+
def __init__(self, to_wrap):
148+
self.to_wrap = to_wrap
149+
self.current_id = None
150+
"""Optional[bytes]: The current transaction ID."""
151+
self.retry_id = None
152+
"""Optional[bytes]: The ID of the first attempted transaction."""
153+
154+
def _reset(self):
155+
"""Unset the transaction IDs."""
156+
self.current_id = None
157+
self.retry_id = None
158+
159+
def _pre_commit(self, transaction, *args, **kwargs):
160+
raise NotImplementedError
161+
162+
def _maybe_commit(self, transaction):
163+
raise NotImplementedError
164+
165+
def __call__(self, transaction, *args, **kwargs):
166+
raise NotImplementedError

google/cloud/firestore_v1/transaction.py

Lines changed: 18 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -20,31 +20,27 @@
2020

2121
import six
2222

23+
from google.cloud.firestore_v1.base_transaction import (
24+
_BaseTransactional,
25+
BaseTransaction,
26+
MAX_ATTEMPTS,
27+
_CANT_BEGIN,
28+
_CANT_ROLLBACK,
29+
_CANT_COMMIT,
30+
_WRITE_READ_ONLY,
31+
_INITIAL_SLEEP,
32+
_MAX_SLEEP,
33+
_MULTIPLIER,
34+
_EXCEED_ATTEMPTS_TEMPLATE,
35+
)
36+
2337
from google.api_core import exceptions
2438
from google.cloud.firestore_v1 import batch
25-
from google.cloud.firestore_v1 import types
2639
from google.cloud.firestore_v1.document import DocumentReference
2740
from google.cloud.firestore_v1.query import Query
2841

2942

30-
MAX_ATTEMPTS = 5
31-
"""int: Default number of transaction attempts (with retries)."""
32-
_CANT_BEGIN = "The transaction has already begun. Current transaction ID: {!r}."
33-
_MISSING_ID_TEMPLATE = "The transaction has no transaction ID, so it cannot be {}."
34-
_CANT_ROLLBACK = _MISSING_ID_TEMPLATE.format("rolled back")
35-
_CANT_COMMIT = _MISSING_ID_TEMPLATE.format("committed")
36-
_WRITE_READ_ONLY = "Cannot perform write operation in read-only transaction."
37-
_INITIAL_SLEEP = 1.0
38-
"""float: Initial "max" for sleep interval. To be used in :func:`_sleep`."""
39-
_MAX_SLEEP = 30.0
40-
"""float: Eventual "max" sleep time. To be used in :func:`_sleep`."""
41-
_MULTIPLIER = 2.0
42-
"""float: Multiplier for exponential backoff. To be used in :func:`_sleep`."""
43-
_EXCEED_ATTEMPTS_TEMPLATE = "Failed to commit transaction in {:d} attempts."
44-
_CANT_RETRY_READ_ONLY = "Only read-write transactions can be retried."
45-
46-
47-
class Transaction(batch.WriteBatch):
43+
class Transaction(batch.WriteBatch, BaseTransaction):
4844
"""Accumulate read-and-write operations to be sent in a transaction.
4945
5046
Args:
@@ -60,9 +56,7 @@ class Transaction(batch.WriteBatch):
6056

6157
def __init__(self, client, max_attempts=MAX_ATTEMPTS, read_only=False):
6258
super(Transaction, self).__init__(client)
63-
self._max_attempts = max_attempts
64-
self._read_only = read_only
65-
self._id = None
59+
BaseTransaction.__init__(self, max_attempts, read_only)
6660

6761
def _add_write_pbs(self, write_pbs):
6862
"""Add `Write`` protobufs to this transaction.
@@ -79,61 +73,6 @@ def _add_write_pbs(self, write_pbs):
7973

8074
super(Transaction, self)._add_write_pbs(write_pbs)
8175

82-
def _options_protobuf(self, retry_id):
83-
"""Convert the current object to protobuf.
84-
85-
The ``retry_id`` value is used when retrying a transaction that
86-
failed (e.g. due to contention). It is intended to be the "first"
87-
transaction that failed (i.e. if multiple retries are needed).
88-
89-
Args:
90-
retry_id (Union[bytes, NoneType]): Transaction ID of a transaction
91-
to be retried.
92-
93-
Returns:
94-
Optional[google.cloud.firestore_v1.types.TransactionOptions]:
95-
The protobuf ``TransactionOptions`` if ``read_only==True`` or if
96-
there is a transaction ID to be retried, else :data:`None`.
97-
98-
Raises:
99-
ValueError: If ``retry_id`` is not :data:`None` but the
100-
transaction is read-only.
101-
"""
102-
if retry_id is not None:
103-
if self._read_only:
104-
raise ValueError(_CANT_RETRY_READ_ONLY)
105-
106-
return types.TransactionOptions(
107-
read_write=types.TransactionOptions.ReadWrite(
108-
retry_transaction=retry_id
109-
)
110-
)
111-
elif self._read_only:
112-
return types.TransactionOptions(
113-
read_only=types.TransactionOptions.ReadOnly()
114-
)
115-
else:
116-
return None
117-
118-
@property
119-
def in_progress(self):
120-
"""Determine if this transaction has already begun.
121-
122-
Returns:
123-
bool: Indicates if the transaction has started.
124-
"""
125-
return self._id is not None
126-
127-
@property
128-
def id(self):
129-
"""Get the current transaction ID.
130-
131-
Returns:
132-
Optional[bytes]: The transaction ID (or :data:`None` if the
133-
current transaction is not in progress).
134-
"""
135-
return self._id
136-
13776
def _begin(self, retry_id=None):
13877
"""Begin the transaction.
13978
@@ -157,14 +96,6 @@ def _begin(self, retry_id=None):
15796
)
15897
self._id = transaction_response.transaction
15998

160-
def _clean_up(self):
161-
"""Clean up the instance after :meth:`_rollback`` or :meth:`_commit``.
162-
163-
This intended to occur on success or failure of the associated RPCs.
164-
"""
165-
self._write_pbs = []
166-
self._id = None
167-
16899
def _rollback(self):
169100
"""Roll back the transaction.
170101
@@ -238,7 +169,7 @@ def get(self, ref_or_query):
238169
)
239170

240171

241-
class _Transactional(object):
172+
class _Transactional(_BaseTransactional):
242173
"""Provide a callable object to use as a transactional decorater.
243174
244175
This is surfaced via
@@ -250,16 +181,7 @@ class _Transactional(object):
250181
"""
251182

252183
def __init__(self, to_wrap):
253-
self.to_wrap = to_wrap
254-
self.current_id = None
255-
"""Optional[bytes]: The current transaction ID."""
256-
self.retry_id = None
257-
"""Optional[bytes]: The ID of the first attempted transaction."""
258-
259-
def _reset(self):
260-
"""Unset the transaction IDs."""
261-
self.current_id = None
262-
self.retry_id = None
184+
super(_Transactional, self).__init__(to_wrap)
263185

264186
def _pre_commit(self, transaction, *args, **kwargs):
265187
"""Begin transaction and call the wrapped callable.

0 commit comments

Comments
 (0)