killbill-aplcache

util: harden context validation in EntitySqlDaoWrapperInvocationHandler Signed-off-by:

2/4/2019 7:52:01 AM

Details

diff --git a/account/src/test/java/org/killbill/billing/account/AccountTestSuiteNoDB.java b/account/src/test/java/org/killbill/billing/account/AccountTestSuiteNoDB.java
index 76beffd..cf1f2df 100644
--- a/account/src/test/java/org/killbill/billing/account/AccountTestSuiteNoDB.java
+++ b/account/src/test/java/org/killbill/billing/account/AccountTestSuiteNoDB.java
@@ -85,6 +85,10 @@ public abstract class AccountTestSuiteNoDB extends GuicyKillbillTestSuiteNoDB {
 
     @AfterMethod(groups = "fast")
     public void afterMethod() throws Exception {
+        if (hasFailed()) {
+            return;
+        }
+
         bus.stop();
     }
 }
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 fd34dee..b09eec7 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
@@ -527,23 +527,27 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
                               final InternalCallContext contextMaybeWithoutAccountRecordId) {
         final TableName destinationTableName = MoreObjects.firstNonNull(tableName.getHistoryTableName(), tableName);
 
-        InternalCallContext context = null;
+        final InternalCallContext context;
+        if (TableName.ACCOUNT.equals(tableName) && ChangeType.INSERT.equals(changeType)) {
+            Preconditions.checkState(entityModelDaoAndHistoryRecordIds.size() == 1, "Bulk insert of accounts isn't supported");
+            final M entityModelDao = Iterables.<M>getFirst(entityModelDaoAndHistoryRecordIds.keySet(), null);
+            // AccountModelDao in practice
+            final TimeZoneAwareEntity accountModelDao = (TimeZoneAwareEntity) entityModelDao;
+            context = internalCallContextFactory.createInternalCallContext(accountModelDao, entityModelDao.getRecordId(), contextMaybeWithoutAccountRecordId);
+        } else if (contextMaybeWithoutAccountRecordId.getAccountRecordId() == null) {
+            Preconditions.checkState(tableName == TableName.TENANT || tableName == TableName.TENANT_BROADCASTS || tableName == TableName.TENANT_KVS || tableName == TableName.TAG_DEFINITIONS || tableName == TableName.SERVICE_BRODCASTS || tableName == TableName.NODE_INFOS,
+                                     "accountRecordId should be set for tableName=%s and changeType=%s", tableName, changeType);
+            context = contextMaybeWithoutAccountRecordId;
+
+        } else {
+            context = contextMaybeWithoutAccountRecordId;
+        }
+
         final Collection<EntityAudit> audits = new LinkedList<EntityAudit>();
         for (final M entityModelDao : entityModelDaoAndHistoryRecordIds.keySet()) {
             final Long targetRecordId = entityModelDaoAndHistoryRecordIds.get(entityModelDao);
-            final EntityAudit audit = new EntityAudit(destinationTableName, targetRecordId, changeType, contextMaybeWithoutAccountRecordId.getCreatedDate());
+            final EntityAudit audit = new EntityAudit(destinationTableName, targetRecordId, changeType, context.getCreatedDate());
             audits.add(audit);
-
-            if (context == null) {
-                // Populate the account record id when creating the account record
-                if (TableName.ACCOUNT.equals(tableName) && ChangeType.INSERT.equals(changeType)) {
-                    // AccountModelDao in practice
-                    final TimeZoneAwareEntity accountModelDao = (TimeZoneAwareEntity) entityModelDao;
-                    context = internalCallContextFactory.createInternalCallContext(accountModelDao, entityModelDao.getRecordId(), contextMaybeWithoutAccountRecordId);
-                } else {
-                    context = contextMaybeWithoutAccountRecordId;
-                }
-            }
         }
 
         sqlDao.insertAuditsFromTransaction(audits, context);