added new_transaction to query and aggregation classes

Author: daniel-sanche

google/cloud/datastore/aggregation.py

@@ -442,13 +442,17 @@ def _next_page(self):
             return None
 
         query_pb = self._build_protobuf()
+        new_transaction_options = None
         transaction = self.client.current_transaction
         if transaction is None:
             transaction_id = None
         else:
             transaction_id = transaction.id
+            if transation._begin_later and transaction._status == transaction._INITIAL:
+                # if transaction hasn't been initialized, initialize it as part of this request
+                new_transaction_options = transaction._options
         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/client.py

@@ -203,11 +203,7 @@ def _extended_lookup(
     new_transaction_options = None
     if transaction is not None:
         transaction_id = transaction.id
-        if (
-            transaction_id is None
-            and transaction._begin_later
-            and transaction._status == transaction._INITIAL
-        ):
+        if and transaction._begin_later and transaction._status == transaction._INITIAL:
             # if transaction hasn't been initialized, initialize it as part of this request
             new_transaction_options = transaction._options
 

google/cloud/datastore/query.py

@@ -778,13 +778,16 @@ def _next_page(self):
             return None
 
         query_pb = self._build_protobuf()
+        new_transaction_options = None
         transaction = self.client.current_transaction
         if transaction is None:
             transaction_id = None
         else:
             transaction_id = transaction.id
+            # if transaction hasn't been initialized, initialize it as part of this request
+            new_transaction_options = transaction._options
         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(

tests/unit/test_aggregation.py

@@ -612,6 +612,54 @@ def test_transaction_id_populated(database_id, aggregation_type, aggregation_arg
     assert read_options.transaction == client.current_transaction.id
 
 
+@pytest.mark.parametrize("database_id", [None, "somedb"], indirect=True)
+@pytest.mark.parametrize(
+    "aggregation_type,aggregation_args",
+    [
+        ("count", ()),
+        (
+            "sum",
+            ("appearances",),
+        ),
+        ("avg", ("appearances",)),
+    ],
+)
+def test_transaction_begin_later(database_id, aggregation_type, aggregation_args):
+    """
+    When an aggregation is run in the context of a transaction with begin_later=True,
+    the new_transaction field should be populated in the request read_options.
+    """
+    import mock
+
+    # make a fake begin_later transaction
+    transaction = mock.Mock()
+    transaction.id = None
+    transaction._begin_later = True
+    transaction._state = transaction._INITIAL
+    mock_datastore_api = mock.Mock()
+    mock_gapic = mock_datastore_api.run_aggregation_query
+    mock_gapic.return_value = _make_aggregation_query_response([])
+    client = _Client(
+        None,
+        datastore_api=mock_datastore_api,
+        database=database_id,
+        transaction=transaction,
+    )
+
+    query = _make_query(client)
+    aggregation_query = _make_aggregation_query(client=client, query=query)
+
+    # initiate requested aggregation (ex count, sum, avg)
+    getattr(aggregation_query, aggregation_type)(*aggregation_args)
+    # run mock query
+    list(aggregation_query.fetch())
+    assert mock_gapic.call_count == 1
+    request = mock_gapic.call_args[1]["request"]
+    read_options = request["read_options"]
+    # ensure new_transaction is populated
+    assert read_options.transaction is None
+    assert read_options.new_transaction == transaction._options
+
 class _Client(object):
     def __init__(
         self,

tests/unit/test_query.py

@@ -698,6 +698,48 @@ def test_transaction_id_populated(database_id):
     assert read_options.transaction == client.current_transaction.id
 
 
+@pytest.mark.parametrize("database_id", [None, "somedb"])
+def test_transaction_begin_later(database_id):
+    """
+    When an aggregation is run in the context of a transaction with begin_later=True,
+    the new_transaction field should be populated in the request read_options.
+    """
+    import mock
+
+
+    # make a fake begin_later transaction
+    transaction = mock.Mock()
+    transaction.id = None
+    transaction._begin_later = True
+    transaction._state = transaction._INITIAL
+
+    mock_datastore_api = mock.Mock()
+    mock_gapic = mock_datastore_api.run_query
+
+    more_results_enum = 3  # NO_MORE_RESULTS
+    response_pb = _make_query_response([], b"", more_results_enum, 0)
+    mock_gapic.return_value = response_pb
+
+    client = _Client(
+        None,
+        datastore_api=mock_datastore_api,
+        database=database_id,
+        transaction=transaction,
+    )
+
+    query = _make_query(client)
+    # run mock query
+    list(query.fetch())
+    assert mock_gapic.call_count == 1
+    request = mock_gapic.call_args[1]["request"]
+    read_options = request["read_options"]
+    # ensure new_transaction is populated
+    assert read_options.transaction is None
+    assert read_options.new_transaction == transaction._options
+
+
+
+
 def test_iterator_constructor_defaults():
     query = object()
     client = object()