killbill-memoizeit

diff --git a/account/pom.xml b/account/pom.xml
index 248888c..409da82 100644
--- a/account/pom.xml
+++ b/account/pom.xml
@@ -49,6 +49,11 @@
             <scope>provided</scope>
         </dependency>
         <dependency>
+            <groupId>com.jayway.awaitility</groupId>
+            <artifactId>awaitility</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
             <groupId>javax.inject</groupId>
             <artifactId>javax.inject</artifactId>
             <scope>provided</scope>

diff --git a/account/src/main/java/org/killbill/billing/account/dao/DefaultAccountDao.java b/account/src/main/java/org/killbill/billing/account/dao/DefaultAccountDao.java
index 1c6c201..b14c677 100644
--- a/account/src/main/java/org/killbill/billing/account/dao/DefaultAccountDao.java
+++ b/account/src/main/java/org/killbill/billing/account/dao/DefaultAccountDao.java
@@ -86,7 +86,7 @@ public class DefaultAccountDao extends EntityDaoBase<AccountModelDao, Account, A
         // We need to re-hydrate the callcontext with the account record id
         final InternalCallContext rehydratedContext = internalCallContextFactory.createInternalCallContext(recordId, context);
         final AccountCreationInternalEvent creationEvent = new DefaultAccountCreationEvent(new DefaultAccountData(savedAccount), savedAccount.getId(),
-                                                                                           context.getAccountRecordId(), context.getTenantRecordId(), context.getUserToken());
+                                                                                           rehydratedContext.getAccountRecordId(), rehydratedContext.getTenantRecordId(), rehydratedContext.getUserToken());
         try {
             eventBus.postFromTransaction(creationEvent, entitySqlDaoWrapperFactory.getHandle().getConnection());
         } catch (final EventBusException e) {

diff --git a/account/src/test/java/org/killbill/billing/account/api/user/TestDefaultAccountUserApi.java b/account/src/test/java/org/killbill/billing/account/api/user/TestDefaultAccountUserApi.java
index 829cb91..39bbb72 100644
--- a/account/src/test/java/org/killbill/billing/account/api/user/TestDefaultAccountUserApi.java
+++ b/account/src/test/java/org/killbill/billing/account/api/user/TestDefaultAccountUserApi.java
@@ -1,7 +1,9 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
+ * Copyright 2014-2015 Groupon, Inc
+ * Copyright 2014-2015 The Billing Project, LLC
  *
- * Ning licenses this file to you under the Apache License, version 2.0
+ * The Billing Project licenses this file to you 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:
  *
@@ -16,21 +18,53 @@
 
 package org.killbill.billing.account.api.user;
 
+import java.util.LinkedList;
+import java.util.List;
 import java.util.UUID;
+import java.util.concurrent.Callable;
 
 import org.killbill.billing.account.AccountTestSuiteWithEmbeddedDB;
 import org.killbill.billing.account.api.Account;
+import org.killbill.billing.account.api.AccountData;
 import org.killbill.billing.account.api.DefaultAccount;
 import org.killbill.billing.account.api.DefaultMutableAccountData;
 import org.killbill.billing.account.api.MutableAccountData;
+import org.killbill.billing.account.dao.AccountModelDao;
 import org.killbill.billing.catalog.api.Currency;
+import org.killbill.billing.events.AccountCreationInternalEvent;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import com.google.common.eventbus.Subscribe;
+
+import static com.jayway.awaitility.Awaitility.await;
+import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.killbill.billing.account.AccountTestUtils.createTestAccount;
 
 public class TestDefaultAccountUserApi extends AccountTestSuiteWithEmbeddedDB {
 
+    @Test(groups = "slow", description = "Test Account creation generates an event")
+    public void testBusEvents() throws Exception {
+        final AccountEventHandler eventHandler = new AccountEventHandler();
+        bus.register(eventHandler);
+
+        final AccountModelDao accountModelDao = createTestAccount();
+        final AccountData defaultAccount = new DefaultAccount(accountModelDao);
+        final Account account = accountUserApi.createAccount(defaultAccount, callContext);
+
+        await().atMost(10, SECONDS).until(new Callable<Boolean>() {
+            @Override
+            public Boolean call() throws Exception {
+                return eventHandler.getAccountCreationInternalEvents().size() == 1;
+            }
+        });
+        final AccountCreationInternalEvent accountCreationInternalEvent = eventHandler.getAccountCreationInternalEvents().get(0);
+        Assert.assertEquals(accountCreationInternalEvent.getId(), account.getId());
+        // account_record_id is most likely 1, although, depending on the DB, we cannot be sure
+        Assert.assertNotNull(accountCreationInternalEvent.getSearchKey1());
+        Assert.assertEquals(accountCreationInternalEvent.getSearchKey2(), internalCallContext.getTenantRecordId());
+    }
+
     @Test(groups = "slow", description = "Test Account update with null values")
     public void testShouldBeAbleToPassNullForSomeFieldsToAvoidUpdate() throws Exception {
         final Account account = accountUserApi.createAccount(new DefaultAccount(createTestAccount()), callContext);
@@ -81,4 +115,18 @@ public class TestDefaultAccountUserApi extends AccountTestSuiteWithEmbeddedDB {
 
         accountUserApi.updateAccount(new DefaultAccount(account.getId(), otherAccount), callContext);
     }
+
+    private static final class AccountEventHandler {
+
+        private final List<AccountCreationInternalEvent> accountCreationInternalEvents = new LinkedList<AccountCreationInternalEvent>();
+
+        @Subscribe
+        public void handleAccountCreationInternalEvent(final AccountCreationInternalEvent creationInternalEvent) {
+            this.accountCreationInternalEvents.add(creationInternalEvent);
+        }
+
+        public List<AccountCreationInternalEvent> getAccountCreationInternalEvents() {
+            return accountCreationInternalEvents;
+        }
+    }
 }

diff --git a/util/src/main/java/org/killbill/billing/util/cache/AccountRecordIdCacheLoader.java b/util/src/main/java/org/killbill/billing/util/cache/AccountRecordIdCacheLoader.java
index 87fff3b..352a2d7 100644
--- a/util/src/main/java/org/killbill/billing/util/cache/AccountRecordIdCacheLoader.java
+++ b/util/src/main/java/org/killbill/billing/util/cache/AccountRecordIdCacheLoader.java
@@ -26,6 +26,7 @@ import javax.inject.Singleton;
 import org.killbill.billing.ObjectType;
 import org.killbill.billing.util.cache.Cachable.CacheType;
 import org.killbill.billing.util.dao.NonEntityDao;
+import org.skife.jdbi.v2.Handle;
 
 import net.sf.ehcache.loader.CacheLoader;
 
@@ -46,7 +47,7 @@ public class AccountRecordIdCacheLoader extends BaseIdCacheLoader implements Cac
     }
 
     @Override
-    protected Object doRetrieveOperation(final String rawKey, final ObjectType objectType) {
-        return nonEntityDao.retrieveAccountRecordIdFromObject(UUID.fromString(rawKey), objectType, null);
+    protected Object doRetrieveOperation(final String rawKey, final ObjectType objectType, final Handle handle) {
+        return nonEntityDao.retrieveAccountRecordIdFromObjectInTransaction(UUID.fromString(rawKey), objectType, null, handle);
     }
 }

diff --git a/util/src/main/java/org/killbill/billing/util/cache/BaseIdCacheLoader.java b/util/src/main/java/org/killbill/billing/util/cache/BaseIdCacheLoader.java
index 22ba37d..2385069 100644
--- a/util/src/main/java/org/killbill/billing/util/cache/BaseIdCacheLoader.java
+++ b/util/src/main/java/org/killbill/billing/util/cache/BaseIdCacheLoader.java
@@ -19,8 +19,7 @@ package org.killbill.billing.util.cache;
 
 import org.killbill.billing.ObjectType;
 import org.killbill.billing.util.cache.Cachable.CacheType;
-import org.killbill.billing.util.dao.NonEntityDao;
-import org.skife.jdbi.v2.IDBI;
+import org.skife.jdbi.v2.Handle;
 
 public abstract class BaseIdCacheLoader extends BaseCacheLoader {
 
@@ -31,8 +30,7 @@ public abstract class BaseIdCacheLoader extends BaseCacheLoader {
     @Override
     public abstract CacheType getCacheType();
 
-
-    protected abstract Object doRetrieveOperation(final String rawKey, final ObjectType objectType);
+    protected abstract Object doRetrieveOperation(final String rawKey, final ObjectType objectType, final Handle handle);
 
     @Override
     public Object load(final Object key, final Object argument) {
@@ -47,12 +45,13 @@ public abstract class BaseIdCacheLoader extends BaseCacheLoader {
 
         final String rawKey;
         if (getCacheType().isKeyPrefixedWithTableName()) {
-            String [] parts = ((String) key).split(CacheControllerDispatcher.CACHE_KEY_SEPARATOR);
+            final String[] parts = ((String) key).split(CacheControllerDispatcher.CACHE_KEY_SEPARATOR);
             rawKey = parts[1];
         } else {
             rawKey = (String) key;
         }
         final ObjectType objectType = ((CacheLoaderArgument) argument).getObjectType();
-        return doRetrieveOperation(rawKey, objectType);
+        final Handle handle = ((CacheLoaderArgument) argument).getHandle();
+        return doRetrieveOperation(rawKey, objectType, handle);
     }
 }

diff --git a/util/src/main/java/org/killbill/billing/util/cache/CacheLoaderArgument.java b/util/src/main/java/org/killbill/billing/util/cache/CacheLoaderArgument.java
index 7035fab..688a481 100644
--- a/util/src/main/java/org/killbill/billing/util/cache/CacheLoaderArgument.java
+++ b/util/src/main/java/org/killbill/billing/util/cache/CacheLoaderArgument.java
@@ -20,21 +20,28 @@ import javax.annotation.Nullable;
 
 import org.killbill.billing.ObjectType;
 import org.killbill.billing.callcontext.InternalTenantContext;
+import org.skife.jdbi.v2.Handle;
 
 public class CacheLoaderArgument {
 
     private final ObjectType objectType;
     private final Object[] args;
     private final InternalTenantContext internalTenantContext;
+    private final Handle handle;
 
     public CacheLoaderArgument(final ObjectType objectType) {
         this(objectType, new Object[]{}, null);
     }
 
     public CacheLoaderArgument(final ObjectType objectType, final Object[] args, @Nullable final InternalTenantContext internalTenantContext) {
+        this(objectType, args, internalTenantContext, null);
+    }
+
+    public CacheLoaderArgument(final ObjectType objectType, final Object[] args, @Nullable final InternalTenantContext internalTenantContext, @Nullable final Handle handle) {
         this.objectType = objectType;
         this.args = args;
         this.internalTenantContext = internalTenantContext;
+        this.handle = handle;
     }
 
     public ObjectType getObjectType() {
@@ -48,4 +55,8 @@ public class CacheLoaderArgument {
     public InternalTenantContext getInternalTenantContext() {
         return internalTenantContext;
     }
+
+    public Handle getHandle() {
+        return handle;
+    }
 }

diff --git a/util/src/main/java/org/killbill/billing/util/cache/ObjectIdCacheLoader.java b/util/src/main/java/org/killbill/billing/util/cache/ObjectIdCacheLoader.java
index b4010bf..2fe7889 100644
--- a/util/src/main/java/org/killbill/billing/util/cache/ObjectIdCacheLoader.java
+++ b/util/src/main/java/org/killbill/billing/util/cache/ObjectIdCacheLoader.java
@@ -23,9 +23,8 @@ import javax.inject.Singleton;
 import org.killbill.billing.ObjectType;
 import org.killbill.billing.util.cache.Cachable.CacheType;
 import org.killbill.billing.util.dao.NonEntityDao;
-import org.skife.jdbi.v2.IDBI;
+import org.skife.jdbi.v2.Handle;
 
-import net.sf.ehcache.CacheException;
 import net.sf.ehcache.loader.CacheLoader;
 
 @Singleton
@@ -45,8 +44,8 @@ public class ObjectIdCacheLoader extends BaseIdCacheLoader implements CacheLoade
     }
 
     @Override
-    protected Object doRetrieveOperation(final String rawKey, final ObjectType objectType) {
+    protected Object doRetrieveOperation(final String rawKey, final ObjectType objectType, final Handle handle) {
         final Long recordId = Long.valueOf(rawKey);
-        return nonEntityDao.retrieveIdFromObject(recordId, objectType, null);
+        return nonEntityDao.retrieveIdFromObjectInTransaction(recordId, objectType, null, handle);
     }
 }

diff --git a/util/src/main/java/org/killbill/billing/util/cache/RecordIdCacheLoader.java b/util/src/main/java/org/killbill/billing/util/cache/RecordIdCacheLoader.java
index 889057b..08efe2e 100644
--- a/util/src/main/java/org/killbill/billing/util/cache/RecordIdCacheLoader.java
+++ b/util/src/main/java/org/killbill/billing/util/cache/RecordIdCacheLoader.java
@@ -26,6 +26,7 @@ import javax.inject.Singleton;
 import org.killbill.billing.ObjectType;
 import org.killbill.billing.util.cache.Cachable.CacheType;
 import org.killbill.billing.util.dao.NonEntityDao;
+import org.skife.jdbi.v2.Handle;
 
 import net.sf.ehcache.loader.CacheLoader;
 
@@ -46,7 +47,7 @@ public class RecordIdCacheLoader extends BaseIdCacheLoader implements CacheLoade
     }
 
     @Override
-    protected Object doRetrieveOperation(final String rawKey, final ObjectType objectType) {
-        return nonEntityDao.retrieveRecordIdFromObject(UUID.fromString(rawKey), objectType, null);
+    protected Object doRetrieveOperation(final String rawKey, final ObjectType objectType, final Handle handle) {
+        return nonEntityDao.retrieveRecordIdFromObjectInTransaction(UUID.fromString(rawKey), objectType, null, handle);
     }
 }

diff --git a/util/src/main/java/org/killbill/billing/util/cache/TenantRecordIdCacheLoader.java b/util/src/main/java/org/killbill/billing/util/cache/TenantRecordIdCacheLoader.java
index fd0501b..17a62af 100644
--- a/util/src/main/java/org/killbill/billing/util/cache/TenantRecordIdCacheLoader.java
+++ b/util/src/main/java/org/killbill/billing/util/cache/TenantRecordIdCacheLoader.java
@@ -26,6 +26,7 @@ import javax.inject.Singleton;
 import org.killbill.billing.ObjectType;
 import org.killbill.billing.util.cache.Cachable.CacheType;
 import org.killbill.billing.util.dao.NonEntityDao;
+import org.skife.jdbi.v2.Handle;
 
 import net.sf.ehcache.loader.CacheLoader;
 
@@ -46,7 +47,7 @@ public class TenantRecordIdCacheLoader extends BaseIdCacheLoader implements Cach
     }
 
     @Override
-    protected Object doRetrieveOperation(final String rawKey, final ObjectType objectType) {
-        return nonEntityDao.retrieveTenantRecordIdFromObject(UUID.fromString(rawKey), objectType, null);
+    protected Object doRetrieveOperation(final String rawKey, final ObjectType objectType, final Handle handle) {
+        return nonEntityDao.retrieveTenantRecordIdFromObjectInTransaction(UUID.fromString(rawKey), objectType, null, handle);
     }
 }

diff --git a/util/src/main/java/org/killbill/billing/util/dao/DefaultNonEntityDao.java b/util/src/main/java/org/killbill/billing/util/dao/DefaultNonEntityDao.java
index 595833e..bd662f7 100644
--- a/util/src/main/java/org/killbill/billing/util/dao/DefaultNonEntityDao.java
+++ b/util/src/main/java/org/killbill/billing/util/dao/DefaultNonEntityDao.java
@@ -21,17 +21,19 @@ import java.util.UUID;
 import javax.annotation.Nullable;
 import javax.inject.Inject;
 
-import org.killbill.billing.callcontext.InternalCallContext;
+import org.killbill.billing.ObjectType;
+import org.killbill.billing.util.cache.CacheController;
 import org.killbill.billing.util.cache.CacheControllerDispatcher;
+import org.killbill.billing.util.cache.CacheLoaderArgument;
 import org.killbill.billing.util.callcontext.InternalCallContextFactory;
 import org.killbill.commons.profiling.Profiling;
 import org.killbill.commons.profiling.Profiling.WithProfilingCallback;
 import org.killbill.commons.profiling.ProfilingFeature.ProfilingFeatureType;
+import org.skife.jdbi.v2.Handle;
 import org.skife.jdbi.v2.IDBI;
+import org.skife.jdbi.v2.sqlobject.SqlObjectBuilder;
 
-import org.killbill.billing.ObjectType;
-import org.killbill.billing.util.cache.CacheController;
-import org.killbill.billing.util.cache.CacheLoaderArgument;
+import com.google.common.base.Preconditions;
 
 public class DefaultNonEntityDao implements NonEntityDao {
 
@@ -46,22 +48,43 @@ public class DefaultNonEntityDao implements NonEntityDao {
         this.withCachingRecordId = new WithCaching<Long, UUID>();
     }
 
-
+    @Override
     public Long retrieveRecordIdFromObject(@Nullable final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache) {
+        return retrieveRecordIdFromObjectInTransaction(objectId, objectType, cache, null);
+    }
+
+    public Long retrieveRecordIdFromObjectInTransaction(@Nullable final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache, @Nullable final Handle handle) {
+        if (objectId == null) {
+            return null;
+        }
+
         final TableName tableName = TableName.fromObjectType(objectType);
+        Preconditions.checkNotNull(tableName, "%s is not a valid ObjectType", objectType);
+
         return withCachingObjectId.withCaching(new OperationRetrieval<UUID, Long>() {
             @Override
             public Long doRetrieve(final UUID objectOrRecordId, final ObjectType objectType) {
-                return nonEntitySqlDao.getRecordIdFromObject(objectId.toString(), tableName.getTableName());
+                final NonEntitySqlDao inTransactionNonEntitySqlDao = handle == null ? nonEntitySqlDao : SqlObjectBuilder.attach(handle, NonEntitySqlDao.class);
+                return inTransactionNonEntitySqlDao.getRecordIdFromObject(objectId.toString(), tableName.getTableName());
             }
         }, objectId, objectType, tableName, cache);
     }
 
+    @Override
     public Long retrieveAccountRecordIdFromObject(@Nullable final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache) {
+        return retrieveAccountRecordIdFromObjectInTransaction(objectId, objectType, cache, null);
+    }
+
+    @Override
+    public Long retrieveAccountRecordIdFromObjectInTransaction(@Nullable final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache, @Nullable final Handle handle) {
         final TableName tableName = TableName.fromObjectType(objectType);
+        Preconditions.checkNotNull(tableName, "%s is not a valid ObjectType", objectType);
+
         return withCachingObjectId.withCaching(new OperationRetrieval<UUID, Long>() {
             @Override
             public Long doRetrieve(final UUID objectId, final ObjectType objectType) {
+                final NonEntitySqlDao inTransactionNonEntitySqlDao = handle == null ? nonEntitySqlDao : SqlObjectBuilder.attach(handle, NonEntitySqlDao.class);
+
                 switch (tableName) {
                     case TENANT:
                     case TAG_DEFINITIONS:
@@ -69,27 +92,37 @@ public class DefaultNonEntityDao implements NonEntityDao {
                         return null;
 
                     case ACCOUNT:
-                        return nonEntitySqlDao.getAccountRecordIdFromAccount(objectId.toString());
+                        return inTransactionNonEntitySqlDao.getAccountRecordIdFromAccount(objectId.toString());
 
                     default:
-                        return nonEntitySqlDao.getAccountRecordIdFromObjectOtherThanAccount(objectId.toString(), tableName.getTableName());
+                        return inTransactionNonEntitySqlDao.getAccountRecordIdFromObjectOtherThanAccount(objectId.toString(), tableName.getTableName());
                 }
             }
         }, objectId, objectType, tableName, cache);
     }
 
+    @Override
     public Long retrieveTenantRecordIdFromObject(@Nullable final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache) {
+        return retrieveTenantRecordIdFromObjectInTransaction(objectId, objectType, cache, null);
+    }
+
+    @Override
+    public Long retrieveTenantRecordIdFromObjectInTransaction(@Nullable final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache, @Nullable final Handle handle) {
         final TableName tableName = TableName.fromObjectType(objectType);
+        Preconditions.checkNotNull(tableName, "%s is not a valid ObjectType", objectType);
+
         return withCachingObjectId.withCaching(new OperationRetrieval<UUID, Long>() {
             @Override
             public Long doRetrieve(final UUID objectId, final ObjectType objectType) {
+                final NonEntitySqlDao inTransactionNonEntitySqlDao = handle == null ? nonEntitySqlDao : SqlObjectBuilder.attach(handle, NonEntitySqlDao.class);
+
                 switch (tableName) {
                     case TENANT:
                         // Explicit cast to Long to avoid NPE (unboxing to long)
-                        return objectId == null ? (Long) 0L : nonEntitySqlDao.getTenantRecordIdFromTenant(objectId.toString());
+                        return objectId == null ? (Long) 0L : inTransactionNonEntitySqlDao.getTenantRecordIdFromTenant(objectId.toString());
 
                     default:
-                        return nonEntitySqlDao.getTenantRecordIdFromObjectOtherThanTenant(objectId.toString(), tableName.getTableName());
+                        return inTransactionNonEntitySqlDao.getTenantRecordIdFromObjectOtherThanTenant(objectId.toString(), tableName.getTableName());
                 }
 
             }
@@ -98,19 +131,27 @@ public class DefaultNonEntityDao implements NonEntityDao {
 
     @Override
     public UUID retrieveIdFromObject(final Long recordId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache) {
+        return retrieveIdFromObjectInTransaction(recordId, objectType, cache, null);
+    }
+
+    @Override
+    public UUID retrieveIdFromObjectInTransaction(final Long recordId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache, @Nullable final Handle handle) {
         if (objectType == ObjectType.TENANT && recordId == InternalCallContextFactory.INTERNAL_TENANT_RECORD_ID) {
             return null;
         }
+
         final TableName tableName = TableName.fromObjectType(objectType);
+        Preconditions.checkNotNull(tableName, "%s is not a valid ObjectType", objectType);
+
         return withCachingRecordId.withCaching(new OperationRetrieval<Long, UUID>() {
             @Override
             public UUID doRetrieve(final Long objectOrRecordId, final ObjectType objectType) {
-                return nonEntitySqlDao.getIdFromObject(recordId, tableName.getTableName());
+                final NonEntitySqlDao inTransactionNonEntitySqlDao = handle == null ? nonEntitySqlDao : SqlObjectBuilder.attach(handle, NonEntitySqlDao.class);
+                return inTransactionNonEntitySqlDao.getIdFromObject(recordId, tableName.getTableName());
             }
         }, recordId, objectType, tableName, cache);
     }
 
-
     @Override
     public Long retrieveLastHistoryRecordIdFromTransaction(@Nullable final Long targetRecordId, final TableName tableName, final NonEntitySqlDao transactional) {
         // There is no caching here because the value returned changes as we add more history records, and so we would need some cache invalidation
@@ -122,8 +163,8 @@ public class DefaultNonEntityDao implements NonEntityDao {
         return nonEntitySqlDao.getHistoryTargetRecordId(recordId, tableName.getTableName());
     }
 
-
     private interface OperationRetrieval<TypeIn, TypeOut> {
+
         public TypeOut doRetrieve(final TypeIn objectOrRecordId, final ObjectType objectType);
     }
 
@@ -144,14 +185,14 @@ public class DefaultNonEntityDao implements NonEntityDao {
             }
             final TypeOut result;
             try {
-                result = prof.executeWithProfiling(ProfilingFeatureType.DAO_DETAILS,  "NonEntityDao (type = " +  objectType + ") cache miss", new WithProfilingCallback<TypeOut>() {
+                result = prof.executeWithProfiling(ProfilingFeatureType.DAO_DETAILS, "NonEntityDao (type = " + objectType + ") cache miss", new WithProfilingCallback<TypeOut>() {
                     @Override
                     public <ExceptionType extends Throwable> TypeOut execute() throws ExceptionType {
                         return op.doRetrieve(objectOrRecordId, objectType);
                     }
                 });
                 return result;
-            } catch (Throwable throwable) {
+            } catch (final Throwable throwable) {
                 // This is only because WithProfilingCallback throws a Throwable...
                 throw new RuntimeException(throwable);
             }

diff --git a/util/src/main/java/org/killbill/billing/util/dao/NonEntityDao.java b/util/src/main/java/org/killbill/billing/util/dao/NonEntityDao.java
index 600cb97..72c9cf2 100644
--- a/util/src/main/java/org/killbill/billing/util/dao/NonEntityDao.java
+++ b/util/src/main/java/org/killbill/billing/util/dao/NonEntityDao.java
@@ -24,18 +24,27 @@ import javax.annotation.Nullable;
 
 import org.killbill.billing.ObjectType;
 import org.killbill.billing.util.cache.CacheController;
+import org.skife.jdbi.v2.Handle;
 
 // This should only be used for internal operations (trusted code, not API), because the context will not be validated!
 public interface NonEntityDao {
 
     public Long retrieveRecordIdFromObject(final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache);
 
+    public Long retrieveRecordIdFromObjectInTransaction(final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache, @Nullable final Handle handle);
+
     public Long retrieveAccountRecordIdFromObject(final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache);
 
+    public Long retrieveAccountRecordIdFromObjectInTransaction(final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache, @Nullable final Handle handle);
+
     public Long retrieveTenantRecordIdFromObject(final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache);
 
+    public Long retrieveTenantRecordIdFromObjectInTransaction(final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache, @Nullable final Handle handle);
+
     public UUID retrieveIdFromObject(final Long recordId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache);
 
+    public UUID retrieveIdFromObjectInTransaction(final Long recordId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache, @Nullable final Handle handle);
+
     // This retrieves from the history table the latest record for which targetId matches the one we are passing
     public Long retrieveLastHistoryRecordIdFromTransaction(final Long targetRecordId, final TableName tableName, final NonEntitySqlDao transactional);
 

diff --git a/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java b/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
index b35649c..d128899 100644
--- a/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
+++ b/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
@@ -243,7 +243,7 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
                     return input instanceof InternalTenantContext;
                 }
             }, null);
-            final CacheLoaderArgument cacheLoaderArgument = new CacheLoaderArgument(objectType, args, internalTenantContext);
+            final CacheLoaderArgument cacheLoaderArgument = new CacheLoaderArgument(objectType, args, internalTenantContext, handle);
             return cache.get(cacheKey, cacheLoaderArgument);
         }
         if (result == null) {

diff --git a/util/src/test/java/org/killbill/billing/dao/MockNonEntityDao.java b/util/src/test/java/org/killbill/billing/dao/MockNonEntityDao.java
index d7bfc00..5893704 100644
--- a/util/src/test/java/org/killbill/billing/dao/MockNonEntityDao.java
+++ b/util/src/test/java/org/killbill/billing/dao/MockNonEntityDao.java
@@ -30,6 +30,7 @@ import org.killbill.billing.util.cache.CacheController;
 import org.killbill.billing.util.dao.NonEntityDao;
 import org.killbill.billing.util.dao.NonEntitySqlDao;
 import org.killbill.billing.util.dao.TableName;
+import org.skife.jdbi.v2.Handle;
 
 public class MockNonEntityDao implements NonEntityDao {
 
@@ -50,27 +51,47 @@ public class MockNonEntityDao implements NonEntityDao {
     }
 
     @Override
+    public Long retrieveRecordIdFromObjectInTransaction(final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache, @Nullable final Handle handle) {
+        return null;
+    }
+
+    @Override
     public Long retrieveAccountRecordIdFromObject(final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache) {
         return accountRecordIdMappings.get(objectId);
     }
 
     @Override
+    public Long retrieveAccountRecordIdFromObjectInTransaction(final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache, @Nullable final Handle handle) {
+        return null;
+    }
+
+    @Override
     public Long retrieveTenantRecordIdFromObject(final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache) {
         return tenantRecordIdMappings.get(objectId);
     }
 
     @Override
-    public Long retrieveLastHistoryRecordIdFromTransaction(final Long targetRecordId, final TableName tableName, final NonEntitySqlDao transactional) {
+    public Long retrieveTenantRecordIdFromObjectInTransaction(final UUID objectId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache, @Nullable final Handle handle) {
         return null;
     }
 
     @Override
-    public Long retrieveHistoryTargetRecordId(final Long recordId, final TableName tableName) {
+    public UUID retrieveIdFromObject(final Long recordId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache) {
         return null;
     }
 
     @Override
-    public UUID retrieveIdFromObject(final Long recordId, final ObjectType objectType, final CacheController<Object, Object> cache) {
+    public UUID retrieveIdFromObjectInTransaction(final Long recordId, final ObjectType objectType, @Nullable final CacheController<Object, Object> cache, @Nullable final Handle handle) {
+        return null;
+    }
+
+    @Override
+    public Long retrieveLastHistoryRecordIdFromTransaction(final Long targetRecordId, final TableName tableName, final NonEntitySqlDao transactional) {
+        return null;
+    }
+
+    @Override
+    public Long retrieveHistoryTargetRecordId(final Long recordId, final TableName tableName) {
         return null;
     }
 }