feat: add new_transaction support (#499)

google/cloud/datastore/aggregation.py

@@ -442,13 +442,11 @@ def _next_page(self):
             return None
 
         query_pb = self._build_protobuf()
-        transaction = self.client.current_transaction
-        if transaction is None:
-            transaction_id = None
-        else:
-            transaction_id = transaction.id
+        transaction_id, new_transaction_options = helpers.get_transaction_options(
+            self.client.current_transaction
+        )
         read_options = helpers.get_read_options(
-            self._eventual, transaction_id, self._read_time
+            self._eventual, transaction_id, self._read_time, new_transaction_options
         )
 
         partition_id = entity_pb2.PartitionId(

google/cloud/datastore/batch.py

@@ -192,6 +192,19 @@ def mutations(self):
         """
         return self._mutations
 
+    def _allow_mutations(self) -> bool:
+        """
+        This method is called to see if the batch is in a proper state to allow
+        `put` and `delete` operations.
+
+        the Transaction subclass overrides this method to support
+        the `begin_later` flag.
+
+        :rtype: bool
+        :returns: True if the batch is in a state to allow mutations.
+        """
+        return self._status == self._IN_PROGRESS
+
     def put(self, entity):
         """Remember an entity's state to be saved during :meth:`commit`.
 
@@ -218,7 +231,7 @@ def put(self, entity):
                  progress, if entity has no key assigned, or if the key's
                  ``project`` does not match ours.
         """
-        if self._status != self._IN_PROGRESS:
+        if not self._allow_mutations():
             raise ValueError("Batch must be in progress to put()")
 
         if entity.key is None:
@@ -248,7 +261,7 @@ def delete(self, key):
                  progress, if key is not complete, or if the key's
                  ``project`` does not match ours.
         """
-        if self._status != self._IN_PROGRESS:
+        if not self._allow_mutations():
             raise ValueError("Batch must be in progress to delete()")
 
         if key.is_partial:
@@ -370,10 +383,12 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         try:
-            if exc_type is None:
-                self.commit()
-            else:
-                self.rollback()
+            # commit or rollback if not in terminal state
+            if self._status not in (self._ABORTED, self._FINISHED):
+                if exc_type is None:
+                    self.commit()
+                else:
+                    self.rollback()
         finally:
             self._client._pop_batch()
 

google/cloud/datastore/client.py

@@ -122,7 +122,7 @@ def _extended_lookup(
     missing=None,
     deferred=None,
     eventual=False,
-    transaction_id=None,
+    transaction=None,
     retry=None,
     timeout=None,
     read_time=None,
@@ -158,10 +158,10 @@ def _extended_lookup(
                      consistency.  If True, request ``EVENTUAL`` read
                      consistency.
 
-    :type transaction_id: str
-    :param transaction_id: If passed, make the request in the scope of
-                           the given transaction.  Incompatible with
-                           ``eventual==True`` or ``read_time``.
+    :type transaction: Transaction
+    :param transaction: If passed, make the request in the scope of
+                        the given transaction.  Incompatible with
+                        ``eventual==True`` or ``read_time``.
 
     :type retry: :class:`google.api_core.retry.Retry`
     :param retry:
@@ -177,7 +177,7 @@ def _extended_lookup(
     :type read_time: datetime
     :param read_time:
         (Optional) Read time to use for read consistency. Incompatible with
-        ``eventual==True`` or ``transaction_id``.
+        ``eventual==True`` or ``transaction``.
         This feature is in private preview.
 
     :type database: str
@@ -199,8 +199,14 @@ def _extended_lookup(
 
     results = []
 
+    transaction_id = None
+    transaction_id, new_transaction_options = helpers.get_transaction_options(
+        transaction
+    )
+    read_options = helpers.get_read_options(
+        eventual, transaction_id, read_time, new_transaction_options
+    )
     loop_num = 0
-    read_options = helpers.get_read_options(eventual, transaction_id, read_time)
     while loop_num < _MAX_LOOPS:  # loop against possible deferred.
         loop_num += 1
         request = {
@@ -214,6 +220,10 @@ def _extended_lookup(
             **kwargs,
         )
 
+        # set new transaction id if we just started a transaction
+        if transaction and lookup_response.transaction:
+            transaction._begin_with_id(lookup_response.transaction)
+
         # Accumulate the new results.
         results.extend(result.entity for result in lookup_response.found)
 
@@ -570,7 +580,7 @@ def get_multi(
             eventual=eventual,
             missing=missing,
             deferred=deferred,
-            transaction_id=transaction and transaction.id,
+            transaction=transaction,
             retry=retry,
             timeout=timeout,
             read_time=read_time,

google/cloud/datastore/helpers.py

@@ -230,7 +230,9 @@ def entity_to_protobuf(entity):
     return entity_pb
 
 
-def get_read_options(eventual, transaction_id, read_time=None):
+def get_read_options(
+    eventual, transaction_id, read_time=None, new_transaction_options=None
+):
     """Validate rules for read options, and assign to the request.
 
     Helper method for ``lookup()`` and ``run_query``.
@@ -245,33 +247,55 @@ def get_read_options(eventual, transaction_id, read_time=None):
     :type read_time: datetime
     :param read_time: Read data from the specified time (may be null). This feature is in private preview.
 
+    :type new_transaction_options: :class:`google.cloud.datastore_v1.types.TransactionOptions`
+    :param new_transaction_options: Options for a new transaction.
+
     :rtype: :class:`.datastore_pb2.ReadOptions`
     :returns: The read options corresponding to the inputs.
     :raises: :class:`ValueError` if more than one of ``eventual==True``,
-            ``transaction``, and ``read_time`` is specified.
+            ``transaction_id``, ``read_time``, and ``new_transaction_options`` is specified.
     """
-    if transaction_id is None:
-        if eventual:
-            if read_time is not None:
-                raise ValueError("eventual must be False when read_time is specified")
-            else:
-                return datastore_pb2.ReadOptions(
-                    read_consistency=datastore_pb2.ReadOptions.ReadConsistency.EVENTUAL
-                )
-        else:
-            if read_time is None:
-                return datastore_pb2.ReadOptions()
-            else:
-                read_time_pb = timestamp_pb2.Timestamp()
-                read_time_pb.FromDatetime(read_time)
-                return datastore_pb2.ReadOptions(read_time=read_time_pb)
-    else:
-        if eventual:
-            raise ValueError("eventual must be False when in a transaction")
-        elif read_time is not None:
-            raise ValueError("transaction and read_time are mutual exclusive")
-        else:
-            return datastore_pb2.ReadOptions(transaction=transaction_id)
+    is_set = [
+        bool(x) for x in (eventual, transaction_id, read_time, new_transaction_options)
+    ]
+    if sum(is_set) > 1:
+        raise ValueError(
+            "At most one of eventual, transaction, or read_time is allowed."
+        )
+    new_options = datastore_pb2.ReadOptions()
+    if transaction_id is not None:
+        new_options.transaction = transaction_id
+    if read_time is not None:
+        read_time_pb = timestamp_pb2.Timestamp()
+        read_time_pb.FromDatetime(read_time)
+        new_options.read_time = read_time_pb
+    if new_transaction_options is not None:
+        new_options.new_transaction = new_transaction_options
+    if eventual:
+        new_options.read_consistency = (
+            datastore_pb2.ReadOptions.ReadConsistency.EVENTUAL
+        )
+    return new_options
+
+
+def get_transaction_options(transaction):
+    """
+    Get the transaction_id or new_transaction_options field from an active transaction object,
+    for use in get_read_options
+
+    These are mutually-exclusive fields, so one or both will be None.
+
+    :rtype: Tuple[Optional[bytes], Optional[google.cloud.datastore_v1.types.TransactionOptions]]
+    :returns: The transaction_id and new_transaction_options fields from the transaction object.
+    """
+    transaction_id, new_transaction_options = None, None
+    if transaction is not None:
+        if transaction.id is not None:
+            transaction_id = transaction.id
+        elif transaction._begin_later and transaction._status == transaction._INITIAL:
+            # If the transaction has not yet been begun, we can use the new_transaction_options field.
+            new_transaction_options = transaction._options
+    return transaction_id, new_transaction_options
 
 
 def key_from_protobuf(pb):

google/cloud/datastore/query.py

@@ -778,13 +778,12 @@ def _next_page(self):
             return None
 
         query_pb = self._build_protobuf()
-        transaction = self.client.current_transaction
-        if transaction is None:
-            transaction_id = None
-        else:
-            transaction_id = transaction.id
+        new_transaction_options = None
+        transaction_id, new_transaction_options = helpers.get_transaction_options(
+            self.client.current_transaction
+        )
         read_options = helpers.get_read_options(
-            self._eventual, transaction_id, self._read_time
+            self._eventual, transaction_id, self._read_time, new_transaction_options
         )
 
         partition_id = entity_pb2.PartitionId(

google/cloud/datastore/transaction.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 """Create / interact with Google Cloud Datastore transactions."""
-
 from google.cloud.datastore.batch import Batch
 from google.cloud.datastore_v1.types import TransactionOptions
 from google.protobuf import timestamp_pb2
@@ -149,15 +148,23 @@ class Transaction(Batch):
     :param read_time: (Optional) Time at which the transaction reads entities.
                       Only allowed when ``read_only=True``. This feature is in private preview.
 
+    :type begin_later: bool
+    :param begin_later: (Optional) If True, the transaction will be started
+                        lazily (i.e. when the first RPC is made). If False,
+                        the transaction will be started as soon as the context manager
+                        is entered. `self.begin()` can also be called manually to begin
+                        the transaction at any time. Default is False.
+
     :raises: :class:`ValueError` if read_time is specified when
              ``read_only=False``.
     """
 
     _status = None
 
-    def __init__(self, client, read_only=False, read_time=None):
+    def __init__(self, client, read_only=False, read_time=None, begin_later=False):
         super(Transaction, self).__init__(client)
         self._id = None
+        self._begin_later = begin_later
 
         if read_only:
             if read_time is not None:
@@ -180,8 +187,8 @@ def __init__(self, client, read_only=False, read_time=None):
     def id(self):
         """Getter for the transaction ID.
 
-        :rtype: str
-        :returns: The ID of the current transaction.
+        :rtype: bytes or None
+        :returns: The ID of the current transaction, or None if not started.
         """
         return self._id
 
@@ -240,6 +247,21 @@ def begin(self, retry=None, timeout=None):
             self._status = self._ABORTED
             raise
 
+    def _begin_with_id(self, transaction_id):
+        """
+        Attach newly created transaction to an existing transaction ID.
+
+        This is used when begin_later is True, when the first lookup request
+        associated with this transaction creates a new transaction ID.
+
+        :type transaction_id: bytes
+        :param transaction_id: ID of the transaction to attach to.
+        """
+        if self._status is not self._INITIAL:
+            raise ValueError("Transaction already begun.")
+        self._id = transaction_id
+        self._status = self._IN_PROGRESS
+
     def rollback(self, retry=None, timeout=None):
         """Rolls back the current transaction.
 
@@ -258,6 +280,12 @@ def rollback(self, retry=None, timeout=None):
             Note that if ``retry`` is specified, the timeout applies
             to each individual attempt.
         """
+        # if transaction has not started, abort it
+        if self._status == self._INITIAL:
+            self._status = self._ABORTED
+            self._id = None
+            return None
+
         kwargs = _make_retry_timeout_kwargs(retry, timeout)
 
         try:
@@ -296,6 +324,15 @@ def commit(self, retry=None, timeout=None):
             Note that if ``retry`` is specified, the timeout applies
             to each individual attempt.
         """
+        # if transaction has not begun, either begin now, or abort if empty
+        if self._status == self._INITIAL:
+            if not self._mutations:
+                self._status = self._ABORTED
+                self._id = None
+                return None
+            else:
+                self.begin()
+
         kwargs = _make_retry_timeout_kwargs(retry, timeout)
 
         try:
@@ -321,3 +358,18 @@ def put(self, entity):
             raise RuntimeError("Transaction is read only")
         else:
             super(Transaction, self).put(entity)
+
+    def __enter__(self):
+        if not self._begin_later:
+            self.begin()
+        self._client._push_batch(self)
+        return self
+
+    def _allow_mutations(self):
+        """
+        Mutations can be added to a transaction if it is in IN_PROGRESS state,
+        or if it is in INITIAL state and the begin_later flag is set.
+        """
+        return self._status == self._IN_PROGRESS or (
+            self._begin_later and self._status == self._INITIAL
+        )