killbill-memoizeit

diff --git a/util/src/main/java/org/killbill/billing/util/dao/EntityHistoryModelDao.java b/util/src/main/java/org/killbill/billing/util/dao/EntityHistoryModelDao.java
index 447f132..9fee7a9 100644
--- a/util/src/main/java/org/killbill/billing/util/dao/EntityHistoryModelDao.java
+++ b/util/src/main/java/org/killbill/billing/util/dao/EntityHistoryModelDao.java
@@ -31,16 +31,18 @@ public class EntityHistoryModelDao<M extends EntityModelDao<E>, E extends Entity
     private Long targetRecordId;
     private M entity;
     private ChangeType changeType;
+    private Long historyRecordId;
 
-    public EntityHistoryModelDao(final UUID id, final M src, final Long targetRecordId, final ChangeType type, final DateTime createdDate) {
+    public EntityHistoryModelDao(final UUID id, final M src, final Long targetRecordId, final ChangeType type, final Long historyRecordId, final DateTime createdDate) {
         super(id, createdDate, createdDate);
         this.changeType = type;
         this.targetRecordId = targetRecordId;
         this.entity = src;
+        this.historyRecordId = historyRecordId;
     }
 
-    public EntityHistoryModelDao(final M src, final Long targetRecordId, final ChangeType type, final DateTime createdDate) {
-        this(UUIDs.randomUUID(), src, targetRecordId, type, createdDate);
+    public EntityHistoryModelDao(final M src, final Long targetRecordId, final ChangeType type, final Long historyRecordId, final DateTime createdDate) {
+        this(UUIDs.randomUUID(), src, targetRecordId, type, historyRecordId, createdDate);
     }
 
     public ChangeType getChangeType() {
@@ -66,4 +68,12 @@ public class EntityHistoryModelDao<M extends EntityModelDao<E>, E extends Entity
     public void setChangeType(final ChangeType changeType) {
         this.changeType = changeType;
     }
+
+    public Long getHistoryRecordId() {
+        return historyRecordId;
+    }
+
+    public void setHistoryRecordId(final Long historyRecordId) {
+        this.historyRecordId = historyRecordId;
+    }
 }

diff --git a/util/src/main/java/org/killbill/billing/util/dao/EntityHistoryModelDaoMapper.java b/util/src/main/java/org/killbill/billing/util/dao/EntityHistoryModelDaoMapper.java
index 5ff331d..49fd117 100644
--- a/util/src/main/java/org/killbill/billing/util/dao/EntityHistoryModelDaoMapper.java
+++ b/util/src/main/java/org/killbill/billing/util/dao/EntityHistoryModelDaoMapper.java
@@ -44,6 +44,9 @@ public class EntityHistoryModelDaoMapper<M extends EntityModelDao<E>, E extends 
         final String changeType = r.getString("change_type");
         final DateTime createdDate = getDateTime(r, "created_date");
 
+        // preserve history record id, as it is needed to reference it with audit log
+        final long historyRecordId = r.getLong("history_record_id");
+
         final M entityModelDao = entityMapper.map(index, r, ctx);
 
         // Hack -- remove the id as it is the history id, not the entity id
@@ -55,6 +58,6 @@ public class EntityHistoryModelDaoMapper<M extends EntityModelDao<E>, E extends 
             ((EntityModelDaoBase) entityModelDao).setAccountRecordId(targetRecordId);
         }
 
-        return new EntityHistoryModelDao(id, entityModelDao, targetRecordId, ChangeType.valueOf(changeType), createdDate);
+        return new EntityHistoryModelDao(id, entityModelDao, targetRecordId, ChangeType.valueOf(changeType), historyRecordId, createdDate);
     }
 }

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 907af4b..e5db8e3 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
@@ -504,7 +504,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());
+        final EntityHistoryModelDao<M, E> history = new EntityHistoryModelDao<M, E>(entityModelDao, entityRecordId, changeType, null, clock.getUTCNow());
         final Long recordId = sqlDao.addHistoryFromTransaction(history, context);
         printSQLWarnings();
         return recordId;

diff --git a/util/src/main/resources/org/killbill/billing/util/entity/dao/EntitySqlDao.sql.stg b/util/src/main/resources/org/killbill/billing/util/entity/dao/EntitySqlDao.sql.stg
index 26faf51..5b861e8 100644
--- a/util/src/main/resources/org/killbill/billing/util/entity/dao/EntitySqlDao.sql.stg
+++ b/util/src/main/resources/org/killbill/billing/util/entity/dao/EntitySqlDao.sql.stg
@@ -336,6 +336,7 @@ auditTableValues() ::= <<
 getHistoryForTargetRecordId() ::= <<
 select
   <idField("")>
+, <prefix>record_id as history_record_id
 , <historyTableFields("t.")>
 <accountRecordIdFieldWithComma("t.")>
 <tenantRecordIdFieldWithComma("t.")>