443 lines
16 KiB
Python
443 lines
16 KiB
Python
|
# Copyright 2017 Google LLC All rights reserved.
|
||
|
|
#
|
||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
# you may not use this file except in compliance with the License.
|
||
|
|
# You may obtain a copy of the License at
|
||
|
|
#
|
||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
#
|
||
|
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
# See the License for the specific language governing permissions and
|
||
|
|
# limitations under the License.
|
||
|
|
|
||
|
|
"""Helpers for applying Google Cloud Firestore changes in a transaction."""
|
||
|
|
|
||
|
|
|
||
|
|
import random
|
||
|
|
import time
|
||
|
|
|
||
|
|
import six
|
||
|
|
|
||
|
|
from google.api_core import exceptions
|
||
|
|
from google.cloud.firestore_v1 import batch
|
||
|
|
from google.cloud.firestore_v1 import types
|
||
|
|
from google.cloud.firestore_v1.document import DocumentReference
|
||
|
|
from google.cloud.firestore_v1.query import Query
|
||
|
|
|
||
|
|
|
||
|
|
MAX_ATTEMPTS = 5
|
||
|
|
"""int: Default number of transaction attempts (with retries)."""
|
||
|
|
_CANT_BEGIN = "The transaction has already begun. Current transaction ID: {!r}."
|
||
|
|
_MISSING_ID_TEMPLATE = "The transaction has no transaction ID, so it cannot be {}."
|
||
|
|
_CANT_ROLLBACK = _MISSING_ID_TEMPLATE.format("rolled back")
|
||
|
|
_CANT_COMMIT = _MISSING_ID_TEMPLATE.format("committed")
|
||
|
|
_WRITE_READ_ONLY = "Cannot perform write operation in read-only transaction."
|
||
|
|
_INITIAL_SLEEP = 1.0
|
||
|
|
"""float: Initial "max" for sleep interval. To be used in :func:`_sleep`."""
|
||
|
|
_MAX_SLEEP = 30.0
|
||
|
|
"""float: Eventual "max" sleep time. To be used in :func:`_sleep`."""
|
||
|
|
_MULTIPLIER = 2.0
|
||
|
|
"""float: Multiplier for exponential backoff. To be used in :func:`_sleep`."""
|
||
|
|
_EXCEED_ATTEMPTS_TEMPLATE = "Failed to commit transaction in {:d} attempts."
|
||
|
|
_CANT_RETRY_READ_ONLY = "Only read-write transactions can be retried."
|
||
|
|
|
||
|
|
|
||
|
|
class Transaction(batch.WriteBatch):
|
||
|
|
"""Accumulate read-and-write operations to be sent in a transaction.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
client (:class:`~google.cloud.firestore_v1.client.Client`):
|
||
|
|
The client that created this transaction.
|
||
|
|
max_attempts (Optional[int]): The maximum number of attempts for
|
||
|
|
the transaction (i.e. allowing retries). Defaults to
|
||
|
|
:attr:`~google.cloud.firestore_v1.transaction.MAX_ATTEMPTS`.
|
||
|
|
read_only (Optional[bool]): Flag indicating if the transaction
|
||
|
|
should be read-only or should allow writes. Defaults to
|
||
|
|
:data:`False`.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, client, max_attempts=MAX_ATTEMPTS, read_only=False):
|
||
|
|
super(Transaction, self).__init__(client)
|
||
|
|
self._max_attempts = max_attempts
|
||
|
|
self._read_only = read_only
|
||
|
|
self._id = None
|
||
|
|
|
||
|
|
def _add_write_pbs(self, write_pbs):
|
||
|
|
"""Add `Write`` protobufs to this transaction.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
write_pbs (List[google.cloud.proto.firestore.v1.\
|
||
|
|
write_pb2.Write]): A list of write protobufs to be added.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
ValueError: If this transaction is read-only.
|
||
|
|
"""
|
||
|
|
if self._read_only:
|
||
|
|
raise ValueError(_WRITE_READ_ONLY)
|
||
|
|
|
||
|
|
super(Transaction, self)._add_write_pbs(write_pbs)
|
||
|
|
|
||
|
|
def _options_protobuf(self, retry_id):
|
||
|
|
"""Convert the current object to protobuf.
|
||
|
|
|
||
|
|
The ``retry_id`` value is used when retrying a transaction that
|
||
|
|
failed (e.g. due to contention). It is intended to be the "first"
|
||
|
|
transaction that failed (i.e. if multiple retries are needed).
|
||
|
|
|
||
|
|
Args:
|
||
|
|
retry_id (Union[bytes, NoneType]): Transaction ID of a transaction
|
||
|
|
to be retried.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Optional[google.cloud.firestore_v1.types.TransactionOptions]:
|
||
|
|
The protobuf ``TransactionOptions`` if ``read_only==True`` or if
|
||
|
|
there is a transaction ID to be retried, else :data:`None`.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
ValueError: If ``retry_id`` is not :data:`None` but the
|
||
|
|
transaction is read-only.
|
||
|
|
"""
|
||
|
|
if retry_id is not None:
|
||
|
|
if self._read_only:
|
||
|
|
raise ValueError(_CANT_RETRY_READ_ONLY)
|
||
|
|
|
||
|
|
return types.TransactionOptions(
|
||
|
|
read_write=types.TransactionOptions.ReadWrite(
|
||
|
|
retry_transaction=retry_id
|
||
|
|
)
|
||
|
|
)
|
||
|
|
elif self._read_only:
|
||
|
|
return types.TransactionOptions(
|
||
|
|
read_only=types.TransactionOptions.ReadOnly()
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
return None
|
||
|
|
|
||
|
|
@property
|
||
|
|
def in_progress(self):
|
||
|
|
"""Determine if this transaction has already begun.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
bool: Indicates if the transaction has started.
|
||
|
|
"""
|
||
|
|
return self._id is not None
|
||
|
|
|
||
|
|
@property
|
||
|
|
def id(self):
|
||
|
|
"""Get the current transaction ID.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Optional[bytes]: The transaction ID (or :data:`None` if the
|
||
|
|
current transaction is not in progress).
|
||
|
|
"""
|
||
|
|
return self._id
|
||
|
|
|
||
|
|
def _begin(self, retry_id=None):
|
||
|
|
"""Begin the transaction.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
retry_id (Optional[bytes]): Transaction ID of a transaction to be
|
||
|
|
retried.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
ValueError: If the current transaction has already begun.
|
||
|
|
"""
|
||
|
|
if self.in_progress:
|
||
|
|
msg = _CANT_BEGIN.format(self._id)
|
||
|
|
raise ValueError(msg)
|
||
|
|
|
||
|
|
transaction_response = self._client._firestore_api.begin_transaction(
|
||
|
|
self._client._database_string,
|
||
|
|
options_=self._options_protobuf(retry_id),
|
||
|
|
metadata=self._client._rpc_metadata,
|
||
|
|
)
|
||
|
|
self._id = transaction_response.transaction
|
||
|
|
|
||
|
|
def _clean_up(self):
|
||
|
|
"""Clean up the instance after :meth:`_rollback`` or :meth:`_commit``.
|
||
|
|
|
||
|
|
This intended to occur on success or failure of the associated RPCs.
|
||
|
|
"""
|
||
|
|
self._write_pbs = []
|
||
|
|
self._id = None
|
||
|
|
|
||
|
|
def _rollback(self):
|
||
|
|
"""Roll back the transaction.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
ValueError: If no transaction is in progress.
|
||
|
|
"""
|
||
|
|
if not self.in_progress:
|
||
|
|
raise ValueError(_CANT_ROLLBACK)
|
||
|
|
|
||
|
|
try:
|
||
|
|
# NOTE: The response is just ``google.protobuf.Empty``.
|
||
|
|
self._client._firestore_api.rollback(
|
||
|
|
self._client._database_string,
|
||
|
|
self._id,
|
||
|
|
metadata=self._client._rpc_metadata,
|
||
|
|
)
|
||
|
|
finally:
|
||
|
|
self._clean_up()
|
||
|
|
|
||
|
|
def _commit(self):
|
||
|
|
"""Transactionally commit the changes accumulated.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
List[:class:`google.cloud.proto.firestore.v1.write_pb2.WriteResult`, ...]:
|
||
|
|
The write results corresponding to the changes committed, returned
|
||
|
|
in the same order as the changes were applied to this transaction.
|
||
|
|
A write result contains an ``update_time`` field.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
ValueError: If no transaction is in progress.
|
||
|
|
"""
|
||
|
|
if not self.in_progress:
|
||
|
|
raise ValueError(_CANT_COMMIT)
|
||
|
|
|
||
|
|
commit_response = _commit_with_retry(self._client, self._write_pbs, self._id)
|
||
|
|
|
||
|
|
self._clean_up()
|
||
|
|
return list(commit_response.write_results)
|
||
|
|
|
||
|
|
def get_all(self, references):
|
||
|
|
"""Retrieves multiple documents from Firestore.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
references (List[.DocumentReference, ...]): Iterable of document
|
||
|
|
references to be retrieved.
|
||
|
|
|
||
|
|
Yields:
|
||
|
|
.DocumentSnapshot: The next document snapshot that fulfills the
|
||
|
|
query, or :data:`None` if the document does not exist.
|
||
|
|
"""
|
||
|
|
return self._client.get_all(references, transaction=self)
|
||
|
|
|
||
|
|
def get(self, ref_or_query):
|
||
|
|
"""
|
||
|
|
Retrieve a document or a query result from the database.
|
||
|
|
Args:
|
||
|
|
ref_or_query The document references or query object to return.
|
||
|
|
Yields:
|
||
|
|
.DocumentSnapshot: The next document snapshot that fulfills the
|
||
|
|
query, or :data:`None` if the document does not exist.
|
||
|
|
"""
|
||
|
|
if isinstance(ref_or_query, DocumentReference):
|
||
|
|
return self._client.get_all([ref_or_query], transaction=self)
|
||
|
|
elif isinstance(ref_or_query, Query):
|
||
|
|
return ref_or_query.stream(transaction=self)
|
||
|
|
else:
|
||
|
|
raise ValueError(
|
||
|
|
'Value for argument "ref_or_query" must be a DocumentReference or a Query.'
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class _Transactional(object):
|
||
|
|
"""Provide a callable object to use as a transactional decorater.
|
||
|
|
|
||
|
|
This is surfaced via
|
||
|
|
:func:`~google.cloud.firestore_v1.transaction.transactional`.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
to_wrap (Callable[[:class:`~google.cloud.firestore_v1.transaction.Transaction`, ...], Any]):
|
||
|
|
A callable that should be run (and retried) in a transaction.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, to_wrap):
|
||
|
|
self.to_wrap = to_wrap
|
||
|
|
self.current_id = None
|
||
|
|
"""Optional[bytes]: The current transaction ID."""
|
||
|
|
self.retry_id = None
|
||
|
|
"""Optional[bytes]: The ID of the first attempted transaction."""
|
||
|
|
|
||
|
|
def _reset(self):
|
||
|
|
"""Unset the transaction IDs."""
|
||
|
|
self.current_id = None
|
||
|
|
self.retry_id = None
|
||
|
|
|
||
|
|
def _pre_commit(self, transaction, *args, **kwargs):
|
||
|
|
"""Begin transaction and call the wrapped callable.
|
||
|
|
|
||
|
|
If the callable raises an exception, the transaction will be rolled
|
||
|
|
back. If not, the transaction will be "ready" for ``Commit`` (i.e.
|
||
|
|
it will have staged writes).
|
||
|
|
|
||
|
|
Args:
|
||
|
|
transaction
|
||
|
|
(:class:`~google.cloud.firestore_v1.transaction.Transaction`):
|
||
|
|
A transaction to execute the callable within.
|
||
|
|
args (Tuple[Any, ...]): The extra positional arguments to pass
|
||
|
|
along to the wrapped callable.
|
||
|
|
kwargs (Dict[str, Any]): The extra keyword arguments to pass
|
||
|
|
along to the wrapped callable.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Any: result of the wrapped callable.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
Exception: Any failure caused by ``to_wrap``.
|
||
|
|
"""
|
||
|
|
# Force the ``transaction`` to be not "in progress".
|
||
|
|
transaction._clean_up()
|
||
|
|
transaction._begin(retry_id=self.retry_id)
|
||
|
|
|
||
|
|
# Update the stored transaction IDs.
|
||
|
|
self.current_id = transaction._id
|
||
|
|
if self.retry_id is None:
|
||
|
|
self.retry_id = self.current_id
|
||
|
|
try:
|
||
|
|
return self.to_wrap(transaction, *args, **kwargs)
|
||
|
|
except: # noqa
|
||
|
|
# NOTE: If ``rollback`` fails this will lose the information
|
||
|
|
# from the original failure.
|
||
|
|
transaction._rollback()
|
||
|
|
raise
|
||
|
|
|
||
|
|
def _maybe_commit(self, transaction):
|
||
|
|
"""Try to commit the transaction.
|
||
|
|
|
||
|
|
If the transaction is read-write and the ``Commit`` fails with the
|
||
|
|
``ABORTED`` status code, it will be retried. Any other failure will
|
||
|
|
not be caught.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
transaction
|
||
|
|
(:class:`~google.cloud.firestore_v1.transaction.Transaction`):
|
||
|
|
The transaction to be ``Commit``-ed.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
bool: Indicating if the commit succeeded.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
transaction._commit()
|
||
|
|
return True
|
||
|
|
except exceptions.GoogleAPICallError as exc:
|
||
|
|
if transaction._read_only:
|
||
|
|
raise
|
||
|
|
|
||
|
|
if isinstance(exc, exceptions.Aborted):
|
||
|
|
# If a read-write transaction returns ABORTED, retry.
|
||
|
|
return False
|
||
|
|
else:
|
||
|
|
raise
|
||
|
|
|
||
|
|
def __call__(self, transaction, *args, **kwargs):
|
||
|
|
"""Execute the wrapped callable within a transaction.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
transaction
|
||
|
|
(:class:`~google.cloud.firestore_v1.transaction.Transaction`):
|
||
|
|
A transaction to execute the callable within.
|
||
|
|
args (Tuple[Any, ...]): The extra positional arguments to pass
|
||
|
|
along to the wrapped callable.
|
||
|
|
kwargs (Dict[str, Any]): The extra keyword arguments to pass
|
||
|
|
along to the wrapped callable.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Any: The result of the wrapped callable.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
ValueError: If the transaction does not succeed in
|
||
|
|
``max_attempts``.
|
||
|
|
"""
|
||
|
|
self._reset()
|
||
|
|
|
||
|
|
for attempt in six.moves.xrange(transaction._max_attempts):
|
||
|
|
result = self._pre_commit(transaction, *args, **kwargs)
|
||
|
|
succeeded = self._maybe_commit(transaction)
|
||
|
|
if succeeded:
|
||
|
|
return result
|
||
|
|
|
||
|
|
# Subsequent requests will use the failed transaction ID as part of
|
||
|
|
# the ``BeginTransactionRequest`` when restarting this transaction
|
||
|
|
# (via ``options.retry_transaction``). This preserves the "spot in
|
||
|
|
# line" of the transaction, so exponential backoff is not required
|
||
|
|
# in this case.
|
||
|
|
|
||
|
|
transaction._rollback()
|
||
|
|
msg = _EXCEED_ATTEMPTS_TEMPLATE.format(transaction._max_attempts)
|
||
|
|
raise ValueError(msg)
|
||
|
|
|
||
|
|
|
||
|
|
def transactional(to_wrap):
|
||
|
|
"""Decorate a callable so that it runs in a transaction.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
to_wrap
|
||
|
|
(Callable[[:class:`~google.cloud.firestore_v1.transaction.Transaction`, ...], Any]):
|
||
|
|
A callable that should be run (and retried) in a transaction.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Callable[[:class:`~google.cloud.firestore_v1.transaction.Transaction`, ...], Any]:
|
||
|
|
the wrapped callable.
|
||
|
|
"""
|
||
|
|
return _Transactional(to_wrap)
|
||
|
|
|
||
|
|
|
||
|
|
def _commit_with_retry(client, write_pbs, transaction_id):
|
||
|
|
"""Call ``Commit`` on the GAPIC client with retry / sleep.
|
||
|
|
|
||
|
|
Retries the ``Commit`` RPC on Unavailable. Usually this RPC-level
|
||
|
|
retry is handled by the underlying GAPICd client, but in this case it
|
||
|
|
doesn't because ``Commit`` is not always idempotent. But here we know it
|
||
|
|
is "idempotent"-like because it has a transaction ID. We also need to do
|
||
|
|
our own retry to special-case the ``INVALID_ARGUMENT`` error.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
client (:class:`~google.cloud.firestore_v1.client.Client`):
|
||
|
|
A client with GAPIC client and configuration details.
|
||
|
|
write_pbs (List[:class:`google.cloud.proto.firestore.v1.write_pb2.Write`, ...]):
|
||
|
|
A ``Write`` protobuf instance to be committed.
|
||
|
|
transaction_id (bytes):
|
||
|
|
ID of an existing transaction that this commit will run in.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
:class:`google.cloud.firestore_v1.types.CommitResponse`:
|
||
|
|
The protobuf response from ``Commit``.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
~google.api_core.exceptions.GoogleAPICallError: If a non-retryable
|
||
|
|
exception is encountered.
|
||
|
|
"""
|
||
|
|
current_sleep = _INITIAL_SLEEP
|
||
|
|
while True:
|
||
|
|
try:
|
||
|
|
return client._firestore_api.commit(
|
||
|
|
client._database_string,
|
||
|
|
write_pbs,
|
||
|
|
transaction=transaction_id,
|
||
|
|
metadata=client._rpc_metadata,
|
||
|
|
)
|
||
|
|
except exceptions.ServiceUnavailable:
|
||
|
|
# Retry
|
||
|
|
pass
|
||
|
|
|
||
|
|
current_sleep = _sleep(current_sleep)
|
||
|
|
|
||
|
|
|
||
|
|
def _sleep(current_sleep, max_sleep=_MAX_SLEEP, multiplier=_MULTIPLIER):
|
||
|
|
"""Sleep and produce a new sleep time.
|
||
|
|
|
||
|
|
.. _Exponential Backoff And Jitter: https://www.awsarchitectureblog.com/\
|
||
|
|
2015/03/backoff.html
|
||
|
|
|
||
|
|
Select a duration between zero and ``current_sleep``. It might seem
|
||
|
|
counterintuitive to have so much jitter, but
|
||
|
|
`Exponential Backoff And Jitter`_ argues that "full jitter" is
|
||
|
|
the best strategy.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
current_sleep (float): The current "max" for sleep interval.
|
||
|
|
max_sleep (Optional[float]): Eventual "max" sleep time
|
||
|
|
multiplier (Optional[float]): Multiplier for exponential backoff.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
float: Newly doubled ``current_sleep`` or ``max_sleep`` (whichever
|
||
|
|
is smaller)
|
||
|
|
"""
|
||
|
|
actual_sleep = random.uniform(0.0, current_sleep)
|
||
|
|
time.sleep(actual_sleep)
|
||
|
|
return min(multiplier * current_sleep, max_sleep)
|