Pass the transaction to Datastore callback contexts.

PiperOrigin-RevId: 879155390
Change-Id: I4f2bb698cf6a475cff0abe448c10c94ac9f57386

Author: ludoch

api/src/main/java/com/google/appengine/api/datastore/BaseAsyncDatastoreServiceImpl.java

@@ -212,7 +212,7 @@ public Future<Map<Key, Entity>> get(@Nullable Transaction txn, Iterable<Key> key
     // Allocate the Map that will receive the result of the RPC here so that PreGet callbacks can
     // add results.
     Map<Key, Entity> resultMap = new HashMap<Key, Entity>();
-    PreGetContext preGetContext = new PreGetContext(this, keyList, resultMap);
+    PreGetContext preGetContext = new PreGetContext(this, txn, keyList, resultMap);
     datastoreServiceConfig.getDatastoreCallbacks().executePreGetCallbacks(preGetContext);
 
     // Don't fetch anything from datastore that was provided by the preGet hooks.
@@ -222,7 +222,7 @@ public Future<Map<Key, Entity>> get(@Nullable Transaction txn, Iterable<Key> key
     Future<Map<Key, Entity>> result = doBatchGet(txn, Sets.newLinkedHashSet(keyList), resultMap);
 
     // Invoke the user post-get callbacks.
-    return new PostLoadFuture(result, datastoreServiceConfig.getDatastoreCallbacks(), this);
+    return new PostLoadFuture(result, datastoreServiceConfig.getDatastoreCallbacks(), this, txn);
   }
 
   private Future<Entity> wrapSingleGet(final Key key, Future<Map<Key, Entity>> futureEntities) {
@@ -268,7 +268,7 @@ public Future<List<Key>> put(@Nullable Transaction txn, Iterable<Entity> entitie
     // Invoke the pre-put callbacks.
     List<Entity> entityList =
         entities instanceof List ? (List<Entity>) entities : Lists.newArrayList(entities);
-    PutContext prePutContext = new PutContext(this, entityList);
+    PutContext prePutContext = new PutContext(this, txn, entityList);
     datastoreServiceConfig.getDatastoreCallbacks().executePrePutCallbacks(prePutContext);
 
     // Do the datastore put RPC on the remaining entities.
@@ -277,7 +277,7 @@ public Future<List<Key>> put(@Nullable Transaction txn, Iterable<Entity> entitie
     if (txn == null) {
       // We're not in a txn so make sure we execute post-put callbacks when
       // the user asks for the result of the future.
-      PutContext postPutContext = new PutContext(this, entityList);
+      PutContext postPutContext = new PutContext(this, txn, entityList);
       result =
           new PostPutFuture(result, datastoreServiceConfig.getDatastoreCallbacks(), postPutContext);
     } else {
@@ -326,7 +326,7 @@ protected Future<Void> runInternal(Transaction txn) {
   @Override
   public Future<Void> delete(@Nullable Transaction txn, Iterable<Key> keys) {
     List<Key> allKeys = keys instanceof List ? (List<Key>) keys : ImmutableList.copyOf(keys);
-    DeleteContext preDeleteContext = new DeleteContext(this, allKeys);
+    DeleteContext preDeleteContext = new DeleteContext(this, txn, allKeys);
     datastoreServiceConfig.getDatastoreCallbacks().executePreDeleteCallbacks(preDeleteContext);
     // NOTE: We are reusing the user's list here, we can do this because
     // we do not hold on to this list after this function returns.
@@ -339,7 +339,7 @@ public Future<Void> delete(@Nullable Transaction txn, Iterable<Key> keys) {
           new PostDeleteFuture(
               result,
               datastoreServiceConfig.getDatastoreCallbacks(),
-              new DeleteContext(this, allKeys));
+              new DeleteContext(this, txn, allKeys));
     } else {
       // We are in a txn so register the entities that have been deleted with
       // the txn so that we execute the appropriate post delete callbacks when
@@ -434,7 +434,7 @@ public PreparedQuery prepare(Query query) {
   @SuppressWarnings("deprecation")
   @Override
   public PreparedQuery prepare(Transaction txn, Query query) {
-    PreQueryContext context = new PreQueryContext(this, query);
+    PreQueryContext context = new PreQueryContext(this, txn, query);
     datastoreServiceConfig.getDatastoreCallbacks().executePreQueryCallbacks(context);
     // Make sure we get the query off the context in case the interceptor has
     // modified it.

api/src/main/java/com/google/appengine/api/datastore/BaseCallbackContext.java

@@ -30,6 +30,7 @@
  */
 abstract class BaseCallbackContext<T> implements CallbackContext<T> {
   private final CurrentTransactionProvider currentTxnProvider;
+  private final Transaction txn;
 
   /** All elements provided to the operation that triggered the callback. */
   private final List<T> elements;
@@ -41,8 +42,10 @@ abstract class BaseCallbackContext<T> implements CallbackContext<T> {
    * @param currentTxnProvider Provides the current transaction
    * @param elements All elements involved in the operation that triggered the callback.
    */
-  BaseCallbackContext(CurrentTransactionProvider currentTxnProvider, List<T> elements) {
+  BaseCallbackContext(
+      CurrentTransactionProvider currentTxnProvider, Transaction txn, List<T> elements) {
     this.currentTxnProvider = Preconditions.checkNotNull(currentTxnProvider);
+    this.txn = txn;
     this.elements = Collections.unmodifiableList(Preconditions.checkNotNull(elements));
   }
 
@@ -53,7 +56,7 @@ public List<T> getElements() {
 
   @Override
   public Transaction getCurrentTransaction() {
-    return currentTxnProvider.getCurrentTransaction(null);
+    return currentTxnProvider.getCurrentTransaction(txn);
   }
 
   @Override

api/src/main/java/com/google/appengine/api/datastore/BaseQueryResultsSource.java

@@ -291,7 +291,8 @@ private int processQueryResult(
     for (Entity entity : entityList) {
       // Since results are streamed back over time we just create one context
       // per result and execute postLoad callbacks within that context.
-      callbacks.executePostLoadCallbacks(new PostLoadContext(currentTransactionProvider, entity));
+      callbacks.executePostLoadCallbacks(
+          new PostLoadContext(currentTransactionProvider, txn, entity));
     }
     totalResults += entityList.size();
     // If the user has pulled back a large number of results without setting

api/src/main/java/com/google/appengine/api/datastore/DeleteContext.java

@@ -24,8 +24,8 @@
  */
 public final class DeleteContext extends BaseCallbackContext<Key> {
 
-  DeleteContext(CurrentTransactionProvider currentTxnProvider, List<Key> keys) {
-    super(currentTxnProvider, keys);
+  DeleteContext(CurrentTransactionProvider currentTxnProvider, Transaction txn, List<Key> keys) {
+    super(currentTxnProvider, txn, keys);
   }
 
   @Override

api/src/main/java/com/google/appengine/api/datastore/PostLoadContext.java

@@ -31,12 +31,14 @@
  */
 public final class PostLoadContext extends BaseCallbackContext<Entity> {
 
-  PostLoadContext(CurrentTransactionProvider currentTransactionProvider, List<Entity> results) {
-    super(currentTransactionProvider, results);
+  PostLoadContext(
+      CurrentTransactionProvider currentTransactionProvider, Transaction txn, List<Entity> results) {
+    super(currentTransactionProvider, txn, results);
   }
 
-  PostLoadContext(CurrentTransactionProvider currentTransactionProvider, Entity result) {
-    this(currentTransactionProvider, Arrays.asList(result));
+  PostLoadContext(
+      CurrentTransactionProvider currentTransactionProvider, Transaction txn, Entity result) {
+    this(currentTransactionProvider, txn, Arrays.asList(result));
   }
 
   @Override

api/src/main/java/com/google/appengine/api/datastore/PostLoadFuture.java

@@ -26,19 +26,22 @@
  */
 class PostLoadFuture extends PostOpFuture<Map<Key, Entity>> {
   private final CurrentTransactionProvider txnProvider;
+  private final Transaction txn;
 
   PostLoadFuture(
       Future<Map<Key, Entity>> delegate,
       DatastoreCallbacks callbacks,
-      CurrentTransactionProvider txnProvider) {
+      CurrentTransactionProvider txnProvider,
+      Transaction txn) {
     super(delegate, callbacks);
     this.txnProvider = txnProvider;
+    this.txn = txn;
   }
 
   @Override
   void executeCallbacks(Map<Key, Entity> result) {
     PostLoadContext postGetContext =
-        new PostLoadContext(txnProvider, Lists.newArrayList(result.values()));
+        new PostLoadContext(txnProvider, txn, Lists.newArrayList(result.values()));
     datastoreCallbacks.executePostLoadCallbacks(postGetContext);
   }
 }

api/src/main/java/com/google/appengine/api/datastore/PreGetContext.java

@@ -34,9 +34,10 @@ public final class PreGetContext extends BaseCallbackContext<Key> {
 
   PreGetContext(
       CurrentTransactionProvider currentTransactionProvider,
+      Transaction txn,
       List<Key> keys,
       Map<Key, Entity> resultMap) {
-    super(currentTransactionProvider, keys);
+    super(currentTransactionProvider, txn, keys);
     this.resultMap = resultMap;
   }
 

api/src/main/java/com/google/appengine/api/datastore/PreQueryContext.java

@@ -29,10 +29,11 @@
 // in as the results are streamed back.
 public final class PreQueryContext extends BaseCallbackContext<Query> {
 
-  PreQueryContext(CurrentTransactionProvider currentTransactionProvider, Query query) {
+  PreQueryContext(
+      CurrentTransactionProvider currentTransactionProvider, Transaction txn, Query query) {
     // Pass in a copy of the provided query so that hooks can mutate without
     // side effects.
-    super(currentTransactionProvider, Arrays.asList(new Query(query)));
+    super(currentTransactionProvider, txn, Arrays.asList(new Query(query)));
   }
 
   @Override

api/src/main/java/com/google/appengine/api/datastore/PutContext.java

@@ -24,8 +24,8 @@
  */
 public final class PutContext extends BaseCallbackContext<Entity> {
 
-  PutContext(CurrentTransactionProvider currentTxnProvider, List<Entity> entities) {
-    super(currentTxnProvider, entities);
+  PutContext(CurrentTransactionProvider currentTxnProvider, Transaction txn, List<Entity> entities) {
+    super(currentTxnProvider, txn, entities);
   }
 
   @Override

api/src/main/java/com/google/appengine/api/datastore/TransactionImpl.java

@@ -95,9 +95,10 @@ private PostCommitFuture(
 
     @Override
     void executeCallbacks(Void ignoreMe) {
-      PutContext putContext = new PutContext(TransactionImpl.this, putEntities);
+      PutContext putContext = new PutContext(TransactionImpl.this, TransactionImpl.this, putEntities);
       callbacks.executePostPutCallbacks(putContext);
-      DeleteContext deleteContext = new DeleteContext(TransactionImpl.this, deletedKeys);
+      DeleteContext deleteContext =
+          new DeleteContext(TransactionImpl.this, TransactionImpl.this, deletedKeys);
       callbacks.executePostDeleteCallbacks(deleteContext);
     }
   }