killbill-memoizeit

diff --git a/util/src/main/java/org/killbill/billing/util/dao/HistorySqlDao.java b/util/src/main/java/org/killbill/billing/util/dao/HistorySqlDao.java
index 1a748c1..5f6d2d1 100644
--- a/util/src/main/java/org/killbill/billing/util/dao/HistorySqlDao.java
+++ b/util/src/main/java/org/killbill/billing/util/dao/HistorySqlDao.java
@@ -25,6 +25,7 @@ import org.killbill.billing.util.entity.Entity;
 import org.killbill.billing.util.entity.dao.EntityModelDao;
 import org.skife.jdbi.v2.sqlobject.Bind;
 import org.skife.jdbi.v2.sqlobject.BindBean;
+import org.skife.jdbi.v2.sqlobject.GetGeneratedKeys;
 import org.skife.jdbi.v2.sqlobject.SqlQuery;
 import org.skife.jdbi.v2.sqlobject.SqlUpdate;
 
@@ -33,8 +34,8 @@ public interface HistorySqlDao<M extends EntityModelDao<E>, E extends Entity> {
     @SqlQuery
     public List<EntityHistoryModelDao<M, E>> getHistoryForTargetRecordId(@Bind("targetRecordId") final long targetRecordId,
                                                                          @BindBean InternalCallContext context);
-
     @SqlUpdate
-    public void addHistoryFromTransaction(@EntityHistoryBinder EntityHistoryModelDao<M, E> history,
+    @GetGeneratedKeys
+    public Long addHistoryFromTransaction(@EntityHistoryBinder EntityHistoryModelDao<M, E> history,
                                           @BindBean InternalCallContext context);
 }

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 c747bfa..7e4c879 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
@@ -468,13 +468,7 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
 
     private Long insertHistory(final Long entityRecordId, final M entityModelDao, final ChangeType changeType, final InternalCallContext context) {
         final EntityHistoryModelDao<M, E> history = new EntityHistoryModelDao<M, E>(entityModelDao, entityRecordId, changeType, clock.getUTCNow());
-
-        sqlDao.addHistoryFromTransaction(history, context);
-
-        final NonEntitySqlDao transactional = SqlObjectBuilder.attach(handle, NonEntitySqlDao.class);
-
-        /* return transactional.getLastHistoryRecordId(entityRecordId, entityModelDao.getHistoryTableName().getTableName()); */
-        return nonEntityDao.retrieveLastHistoryRecordIdFromTransaction(entityRecordId, entityModelDao.getHistoryTableName(), transactional);
+        return sqlDao.addHistoryFromTransaction(history, context);
     }
 
     private void insertAudits(final TableName tableName, final M entityModelDao, final Long entityRecordId, final Long historyRecordId, final ChangeType changeType, final InternalCallContext contextMaybeWithoutAccountRecordId) {