1# Copyright 2014 Google LLC
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"""Create / interact with Google Cloud Datastore transactions."""
16
17from google.cloud.datastore.batch import Batch
18from google.cloud.datastore_v1.types import TransactionOptions
19
20
21def _make_retry_timeout_kwargs(retry, timeout):
22    """Helper: make optional retry / timeout kwargs dict."""
23    kwargs = {}
24
25    if retry is not None:
26        kwargs["retry"] = retry
27
28    if timeout is not None:
29        kwargs["timeout"] = timeout
30
31    return kwargs
32
33
34class Transaction(Batch):
35    """An abstraction representing datastore Transactions.
36
37    Transactions can be used to build up a bulk mutation and ensure all
38    or none succeed (transactionally).
39
40    For example, the following snippet of code will put the two ``save``
41    operations (either ``insert`` or ``upsert``) into the same
42    mutation, and execute those within a transaction:
43
44    .. testsetup:: txn
45
46        import uuid
47
48        from google.cloud import datastore
49
50        unique = str(uuid.uuid4())[0:8]
51        client = datastore.Client(namespace='ns{}'.format(unique))
52
53    .. doctest:: txn
54
55        >>> entity1 = datastore.Entity(client.key('EntityKind', 1234))
56        >>> entity2 = datastore.Entity(client.key('EntityKind', 2345))
57        >>> with client.transaction():
58        ...     client.put_multi([entity1, entity2])
59
60    Because it derives from :class:`~google.cloud.datastore.batch.Batch`,
61    :class:`Transaction` also provides :meth:`put` and :meth:`delete` methods:
62
63    .. doctest:: txn
64
65       >>> with client.transaction() as xact:
66       ...     xact.put(entity1)
67       ...     xact.delete(entity2.key)
68
69    By default, the transaction is rolled back if the transaction block
70    exits with an error:
71
72    .. doctest:: txn
73
74        >>> def do_some_work():
75        ...    return
76        >>> class SomeException(Exception):
77        ...    pass
78        >>> with client.transaction():
79        ...     do_some_work()
80        ...     raise SomeException  # rolls back
81        Traceback (most recent call last):
82          ...
83        SomeException
84
85    If the transaction block exits without an exception, it will commit
86    by default.
87
88    .. warning::
89
90        Inside a transaction, automatically assigned IDs for
91        entities will not be available at save time!  That means, if you
92        try:
93
94        .. doctest:: txn
95
96            >>> with client.transaction():
97            ...     thing1 = datastore.Entity(key=client.key('Thing'))
98            ...     client.put(thing1)
99
100       ``thing1`` won't have a complete key until the transaction is
101       committed.
102
103       Once you exit the transaction (or call :meth:`commit`), the
104       automatically generated ID will be assigned to the entity:
105
106       .. doctest:: txn
107
108          >>> with client.transaction():
109          ...     thing2 = datastore.Entity(key=client.key('Thing'))
110          ...     client.put(thing2)
111          ...     print(thing2.key.is_partial)  # There is no ID on this key.
112          ...
113          True
114          >>> print(thing2.key.is_partial)  # There *is* an ID.
115          False
116
117    If you don't want to use the context manager you can initialize a
118    transaction manually:
119
120    .. doctest:: txn
121
122       >>> transaction = client.transaction()
123       >>> transaction.begin()
124       >>>
125       >>> thing3 = datastore.Entity(key=client.key('Thing'))
126       >>> transaction.put(thing3)
127       >>>
128       >>> transaction.commit()
129
130    .. testcleanup:: txn
131
132        with client.batch() as batch:
133            batch.delete(client.key('EntityKind', 1234))
134            batch.delete(client.key('EntityKind', 2345))
135            batch.delete(thing1.key)
136            batch.delete(thing2.key)
137            batch.delete(thing3.key)
138
139    :type client: :class:`google.cloud.datastore.client.Client`
140    :param client: the client used to connect to datastore.
141
142    :type read_only: bool
143    :param read_only: indicates the transaction is read only.
144    """
145
146    _status = None
147
148    def __init__(self, client, read_only=False):
149        super(Transaction, self).__init__(client)
150        self._id = None
151
152        if read_only:
153            options = TransactionOptions(read_only=TransactionOptions.ReadOnly())
154        else:
155            options = TransactionOptions()
156
157        self._options = options
158
159    @property
160    def id(self):
161        """Getter for the transaction ID.
162
163        :rtype: str
164        :returns: The ID of the current transaction.
165        """
166        return self._id
167
168    def current(self):
169        """Return the topmost transaction.
170
171        .. note::
172
173            If the topmost element on the stack is not a transaction,
174            returns None.
175
176        :rtype: :class:`google.cloud.datastore.transaction.Transaction` or None
177        :returns: The current transaction (if any are active).
178        """
179        top = super(Transaction, self).current()
180        if isinstance(top, Transaction):
181            return top
182
183    def begin(self, retry=None, timeout=None):
184        """Begins a transaction.
185
186        This method is called automatically when entering a with
187        statement, however it can be called explicitly if you don't want
188        to use a context manager.
189
190        :type retry: :class:`google.api_core.retry.Retry`
191        :param retry:
192            A retry object used to retry requests. If ``None`` is specified,
193            requests will be retried using a default configuration.
194
195        :type timeout: float
196        :param timeout:
197            Time, in seconds, to wait for the request to complete.
198            Note that if ``retry`` is specified, the timeout applies
199            to each individual attempt.
200
201        :raises: :class:`~exceptions.ValueError` if the transaction has
202                 already begun.
203        """
204        super(Transaction, self).begin()
205
206        kwargs = _make_retry_timeout_kwargs(retry, timeout)
207
208        request = {
209            "project_id": self.project,
210            "transaction_options": self._options,
211        }
212        try:
213            response_pb = self._client._datastore_api.begin_transaction(
214                request=request, **kwargs
215            )
216            self._id = response_pb.transaction
217        except:  # noqa: E722 do not use bare except, specify exception instead
218            self._status = self._ABORTED
219            raise
220
221    def rollback(self, retry=None, timeout=None):
222        """Rolls back the current transaction.
223
224        This method has necessary side-effects:
225
226        - Sets the current transaction's ID to None.
227
228        :type retry: :class:`google.api_core.retry.Retry`
229        :param retry:
230            A retry object used to retry requests. If ``None`` is specified,
231            requests will be retried using a default configuration.
232
233        :type timeout: float
234        :param timeout:
235            Time, in seconds, to wait for the request to complete.
236            Note that if ``retry`` is specified, the timeout applies
237            to each individual attempt.
238        """
239        kwargs = _make_retry_timeout_kwargs(retry, timeout)
240
241        try:
242            # No need to use the response it contains nothing.
243            self._client._datastore_api.rollback(
244                request={"project_id": self.project, "transaction": self._id}, **kwargs
245            )
246        finally:
247            super(Transaction, self).rollback()
248            # Clear our own ID in case this gets accidentally reused.
249            self._id = None
250
251    def commit(self, retry=None, timeout=None):
252        """Commits the transaction.
253
254        This is called automatically upon exiting a with statement,
255        however it can be called explicitly if you don't want to use a
256        context manager.
257
258        This method has necessary side-effects:
259
260        - Sets the current transaction's ID to None.
261
262        :type retry: :class:`google.api_core.retry.Retry`
263        :param retry:
264            A retry object used to retry requests. If ``None`` is specified,
265            requests will be retried using a default configuration.
266
267        :type timeout: float
268        :param timeout:
269            Time, in seconds, to wait for the request to complete.
270            Note that if ``retry`` is specified, the timeout applies
271            to each individual attempt.
272        """
273        kwargs = _make_retry_timeout_kwargs(retry, timeout)
274
275        try:
276            super(Transaction, self).commit(**kwargs)
277        finally:
278            # Clear our own ID in case this gets accidentally reused.
279            self._id = None
280
281    def put(self, entity):
282        """Adds an entity to be committed.
283
284        Ensures the transaction is not marked readonly.
285        Please see documentation at
286        :meth:`~google.cloud.datastore.batch.Batch.put`
287
288        :type entity: :class:`~google.cloud.datastore.entity.Entity`
289        :param entity: the entity to be saved.
290
291        :raises: :class:`RuntimeError` if the transaction
292                 is marked ReadOnly
293        """
294        if "read_only" in self._options:
295            raise RuntimeError("Transaction is read only")
296        else:
297            super(Transaction, self).put(entity)
298