killbill-memoizeit

diff --git a/account/src/main/java/com/ning/billing/account/dao/AuditedAccountEmailDao.java b/account/src/main/java/com/ning/billing/account/dao/AuditedAccountEmailDao.java
index 6ea0786..664d9cc 100644
--- a/account/src/main/java/com/ning/billing/account/dao/AuditedAccountEmailDao.java
+++ b/account/src/main/java/com/ning/billing/account/dao/AuditedAccountEmailDao.java
@@ -16,6 +16,13 @@
 
 package com.ning.billing.account.dao;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
+
+import org.skife.jdbi.v2.IDBI;
+import org.skife.jdbi.v2.sqlobject.mixins.Transmogrifier;
+
 import com.google.inject.Inject;
 import com.ning.billing.account.api.AccountEmail;
 import com.ning.billing.util.callcontext.CallContext;
@@ -23,14 +30,8 @@ import com.ning.billing.util.dao.AuditedCollectionDaoBase;
 import com.ning.billing.util.dao.ObjectType;
 import com.ning.billing.util.dao.TableName;
 import com.ning.billing.util.entity.collection.dao.UpdatableEntityCollectionSqlDao;
-import org.skife.jdbi.v2.IDBI;
-import org.skife.jdbi.v2.sqlobject.mixins.Transmogrifier;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.UUID;
-
-public class AuditedAccountEmailDao extends AuditedCollectionDaoBase<AccountEmail> implements AccountEmailDao {
+public class AuditedAccountEmailDao extends AuditedCollectionDaoBase<AccountEmail, AccountEmail> implements AccountEmailDao {
     private final AccountEmailSqlDao accountEmailSqlDao;
 
     @Inject
@@ -39,6 +40,11 @@ public class AuditedAccountEmailDao extends AuditedCollectionDaoBase<AccountEmai
     }
 
     @Override
+    protected AccountEmail getEquivalenceObjectFor(AccountEmail obj) {
+        return obj;
+    }
+
+    @Override
     public List<AccountEmail> getEmails(final UUID accountId) {
         return new ArrayList<AccountEmail>(super.loadEntities(accountId, ObjectType.ACCOUNT_EMAIL).values());
     }
@@ -53,6 +59,7 @@ public class AuditedAccountEmailDao extends AuditedCollectionDaoBase<AccountEmai
         return entity.getEmail();
     }
 
+    @Override
     public void test() {
         accountEmailSqlDao.test();
     }

diff --git a/util/src/main/java/com/ning/billing/util/customfield/dao/AuditedCustomFieldDao.java b/util/src/main/java/com/ning/billing/util/customfield/dao/AuditedCustomFieldDao.java
index 6cf6f0c..e5d5ad0 100644
--- a/util/src/main/java/com/ning/billing/util/customfield/dao/AuditedCustomFieldDao.java
+++ b/util/src/main/java/com/ning/billing/util/customfield/dao/AuditedCustomFieldDao.java
@@ -16,15 +16,16 @@
 
 package com.ning.billing.util.customfield.dao;
 
+import org.skife.jdbi.v2.IDBI;
+import org.skife.jdbi.v2.sqlobject.mixins.Transmogrifier;
+
 import com.google.inject.Inject;
 import com.ning.billing.util.customfield.CustomField;
 import com.ning.billing.util.dao.AuditedCollectionDaoBase;
 import com.ning.billing.util.dao.TableName;
 import com.ning.billing.util.entity.collection.dao.UpdatableEntityCollectionSqlDao;
-import org.skife.jdbi.v2.IDBI;
-import org.skife.jdbi.v2.sqlobject.mixins.Transmogrifier;
 
-public class AuditedCustomFieldDao extends AuditedCollectionDaoBase<CustomField> implements CustomFieldDao {
+public class AuditedCustomFieldDao extends AuditedCollectionDaoBase<CustomField, String> implements CustomFieldDao {
     private final CustomFieldSqlDao dao;
 
     @Inject
@@ -38,6 +39,11 @@ public class AuditedCustomFieldDao extends AuditedCollectionDaoBase<CustomField>
     }
 
     @Override
+    protected String getEquivalenceObjectFor(CustomField obj) {
+        return obj.getName();
+    }
+
+    @Override
     protected UpdatableEntityCollectionSqlDao<CustomField> transmogrifyDao(Transmogrifier transactionalDao) {
         return transactionalDao.become(CustomFieldSqlDao.class);
     }

diff --git a/util/src/main/java/com/ning/billing/util/customfield/StringCustomField.java b/util/src/main/java/com/ning/billing/util/customfield/StringCustomField.java
index dc593db..7b7268a 100644
--- a/util/src/main/java/com/ning/billing/util/customfield/StringCustomField.java
+++ b/util/src/main/java/com/ning/billing/util/customfield/StringCustomField.java
@@ -52,6 +52,11 @@ public class StringCustomField extends UpdatableEntityBase implements CustomFiel
     }
 
     @Override
+    public String toString() {
+        return "StringCustomField [name=" + name + ", value=" + value + ", id=" + id + "]";
+    }
+
+    @Override
     public int hashCode() {
         final int prime = 31;
         int result = 1;

diff --git a/util/src/main/java/com/ning/billing/util/dao/AuditedCollectionDaoBase.java b/util/src/main/java/com/ning/billing/util/dao/AuditedCollectionDaoBase.java
index 8acc42e..3cf6620 100644
--- a/util/src/main/java/com/ning/billing/util/dao/AuditedCollectionDaoBase.java
+++ b/util/src/main/java/com/ning/billing/util/dao/AuditedCollectionDaoBase.java
@@ -19,7 +19,6 @@ package com.ning.billing.util.dao;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -27,53 +26,65 @@ import java.util.UUID;
 
 import org.skife.jdbi.v2.sqlobject.mixins.Transmogrifier;
 
-import com.google.common.base.Function;
-import com.google.common.base.Predicate;
-import com.google.common.collect.Collections2;
 import com.google.common.collect.Sets;
 import com.ning.billing.util.ChangeType;
 import com.ning.billing.util.callcontext.CallContext;
 import com.ning.billing.util.entity.Entity;
 import com.ning.billing.util.entity.collection.dao.UpdatableEntityCollectionSqlDao;
 
-public abstract class AuditedCollectionDaoBase<T extends Entity> implements AuditedCollectionDao<T> {
-    private static final class IdInSetPredicate<T extends Entity> implements Predicate<T> {
-        private final Set<UUID> ids;
-
-        public IdInSetPredicate(Set<UUID> ids) {
-            this.ids = ids;
-        }
-
-        @Override
-        public boolean apply(T entity) {
-            return ids.contains(entity.getId());
-        }
-    }
-
-    private final Function<T, UUID> entityIdExtractor = new Function<T, UUID>() {
-        @Override
-        public UUID apply(T entity) {
-            return entity.getId();
-        }
-    };
+public abstract class AuditedCollectionDaoBase<T extends Entity, V> implements AuditedCollectionDao<T> {
+    /**
+     * Returns equivalence object for the entities, so that dao
+     * can figure out if entities have changed (UPDATE statement) or
+     * are new (INSERT statement).
+     * If two entities return the equivalence objects that are equal themselves
+     * (and have the same hashCode), then the entities are equivalent.
+     * For example, two custom field instances are equivalent (describe the same
+     * custom field) if their name is the same (the equivalence object is the
+     * name string). The instances can still have different custom field values
+     * (represent two different 'assignments' to that
+     * field), which should result in UPDATE statements in the dao.
+     */
+    protected abstract V getEquivalenceObjectFor(T obj);
 
     @Override
-    public void saveEntitiesFromTransaction(Transmogrifier transactionalDao, UUID objectId, ObjectType objectType, List<T> entities, CallContext context) {
+    public void saveEntitiesFromTransaction(Transmogrifier transactionalDao, UUID objectId, ObjectType objectType, List<T> newEntities, CallContext context) {
         UpdatableEntityCollectionSqlDao<T> dao = transmogrifyDao(transactionalDao);
 
-        List<T> existingEntities = dao.load(objectId.toString(), objectType);
-
         // get list of existing entities
-        Set<UUID> currentObjIds = new HashSet<UUID>(Collections2.transform(existingEntities, entityIdExtractor));
-        Set<UUID> updatedObjIds = new HashSet<UUID>(Collections2.transform(entities, entityIdExtractor));
+        List<T> currentEntities = dao.load(objectId.toString(), objectType);
 
-        Set<UUID> idsOfObjsToRemove = Sets.difference(currentObjIds, updatedObjIds);
-        Set<UUID> idsOfObjsToAdd = Sets.difference(updatedObjIds, currentObjIds);
-        Set<UUID> idsOfObjsToUpdate = Sets.intersection(currentObjIds, updatedObjIds);
+        Map<V, T> currentObjs = new HashMap<V, T>(currentEntities.size());
+        Map<V, T> updatedObjs = new HashMap<V, T>(newEntities.size());
 
-        Collection<T> objsToRemove = Collections2.filter(existingEntities, new IdInSetPredicate<T>(idsOfObjsToRemove));
-        Collection<T> objsToAdd = Collections2.filter(entities, new IdInSetPredicate<T>(idsOfObjsToAdd));
-        Collection<T> objsToUpdate = Collections2.filter(existingEntities, new IdInSetPredicate<T>(idsOfObjsToUpdate));
+        for (T currentObj : currentEntities) {
+            currentObjs.put(getEquivalenceObjectFor(currentObj), currentObj);
+        }
+        for (T updatedObj : newEntities) {
+            updatedObjs.put(getEquivalenceObjectFor(updatedObj), updatedObj);
+        }
+
+        Set<V> equivToRemove = Sets.difference(currentObjs.keySet(), updatedObjs.keySet());
+        Set<V> equivToAdd = Sets.difference(updatedObjs.keySet(), currentObjs.keySet());
+        Set<V> equivToCheckForUpdate = Sets.intersection(updatedObjs.keySet(), currentObjs.keySet());
+
+        List<T> objsToAdd = new ArrayList<T>(equivToAdd.size());
+        List<T> objsToRemove = new ArrayList<T>(equivToRemove.size());
+        List<T> objsToUpdate = new ArrayList<T>(equivToCheckForUpdate.size());
+
+        for (V equiv : equivToAdd) {
+            objsToAdd.add(updatedObjs.get(equiv));
+        }
+        for (V equiv : equivToRemove) {
+            objsToRemove.add(currentObjs.get(equiv));
+        }
+        for (V equiv : equivToCheckForUpdate) {
+            T currentObj = currentObjs.get(equiv);
+            T updatedObj = updatedObjs.get(equiv);
+            if (!currentObj.equals(updatedObj)) {
+                objsToUpdate.add(updatedObj);
+            }
+        }
 
         if (objsToAdd.size() != 0) {
             dao.insertFromTransaction(objectId.toString(), objectType, objsToAdd, context);

diff --git a/util/src/main/java/com/ning/billing/util/dao/AuditSqlDao.java b/util/src/main/java/com/ning/billing/util/dao/AuditSqlDao.java
index d488c6b..4f9c256 100644
--- a/util/src/main/java/com/ning/billing/util/dao/AuditSqlDao.java
+++ b/util/src/main/java/com/ning/billing/util/dao/AuditSqlDao.java
@@ -16,15 +16,16 @@
 
 package com.ning.billing.util.dao;
 
-import com.ning.billing.util.callcontext.CallContext;
-import com.ning.billing.util.callcontext.CallContextBinder;
+import java.util.List;
+
 import org.skife.jdbi.v2.sqlobject.Bind;
 import org.skife.jdbi.v2.sqlobject.SqlBatch;
 import org.skife.jdbi.v2.sqlobject.SqlQuery;
 import org.skife.jdbi.v2.sqlobject.SqlUpdate;
 import org.skife.jdbi.v2.sqlobject.stringtemplate.ExternalizedSqlViaStringTemplate3;
 
-import java.util.List;
+import com.ning.billing.util.callcontext.CallContext;
+import com.ning.billing.util.callcontext.CallContextBinder;
 
 @ExternalizedSqlViaStringTemplate3
 public interface AuditSqlDao {
@@ -41,4 +42,5 @@ public interface AuditSqlDao {
 
     @SqlQuery
     public Long getHistoryRecordId(@Bind("recordId") final Long recordId);
+
 }

diff --git a/util/src/main/java/com/ning/billing/util/tag/dao/AuditedTagDao.java b/util/src/main/java/com/ning/billing/util/tag/dao/AuditedTagDao.java
index 8118788..a803660 100644
--- a/util/src/main/java/com/ning/billing/util/tag/dao/AuditedTagDao.java
+++ b/util/src/main/java/com/ning/billing/util/tag/dao/AuditedTagDao.java
@@ -45,7 +45,7 @@ import com.ning.billing.util.tag.DescriptiveTag;
 import com.ning.billing.util.tag.Tag;
 import com.ning.billing.util.tag.TagDefinition;
 
-public class AuditedTagDao extends AuditedCollectionDaoBase<Tag> implements TagDao {
+public class AuditedTagDao extends AuditedCollectionDaoBase<Tag, Tag> implements TagDao {
     private final TagSqlDao tagSqlDao;
 
     @Inject
@@ -54,6 +54,11 @@ public class AuditedTagDao extends AuditedCollectionDaoBase<Tag> implements TagD
     }
 
     @Override
+    protected Tag getEquivalenceObjectFor(Tag obj) {
+        return obj;
+    }
+
+    @Override
     public void insertTag(final UUID objectId, final ObjectType objectType,
                           final TagDefinition tagDefinition, final CallContext context) {
         tagSqlDao.inTransaction(new Transaction<Void, TagSqlDao>() {

diff --git a/util/src/main/java/com/ning/billing/util/tag/DefaultControlTag.java b/util/src/main/java/com/ning/billing/util/tag/DefaultControlTag.java
index 613f043..330b204 100644
--- a/util/src/main/java/com/ning/billing/util/tag/DefaultControlTag.java
+++ b/util/src/main/java/com/ning/billing/util/tag/DefaultControlTag.java
@@ -40,7 +40,7 @@ public class DefaultControlTag extends DescriptiveTag implements ControlTag {
 
     @Override
     public String toString() {
-        return controlTagType.toString();
+        return "DefaultControlTag [controlTagType=" + controlTagType + ", id=" + id + "]";
     }
 
     @Override

diff --git a/util/src/main/java/com/ning/billing/util/tag/DescriptiveTag.java b/util/src/main/java/com/ning/billing/util/tag/DescriptiveTag.java
index c3d8f93..3f875be 100644
--- a/util/src/main/java/com/ning/billing/util/tag/DescriptiveTag.java
+++ b/util/src/main/java/com/ning/billing/util/tag/DescriptiveTag.java
@@ -48,24 +48,34 @@ public class DescriptiveTag extends EntityBase implements Tag {
 
     @Override
     public String toString() {
-        return tagDefinitionName;
+        return "DescriptiveTag [tagDefinitionName=" + tagDefinitionName + ", id=" + id + "]";
     }
 
     @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-
-        DescriptiveTag that = (DescriptiveTag) o;
+    public int hashCode() {
+        final int prime = 31;
+        int result = 1;
+        result = prime * result + ((tagDefinitionName == null) ? 0
+                                                              : tagDefinitionName.hashCode());
+        return result;
+    }
 
-        if (tagDefinitionName != null ? !tagDefinitionName.equals(that.tagDefinitionName) : that.tagDefinitionName != null)
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj)
+            return true;
+        if (obj == null)
+            return false;
+        if (getClass() != obj.getClass())
+            return false;
+        DescriptiveTag other = (DescriptiveTag) obj;
+        if (tagDefinitionName == null) {
+            if (other.tagDefinitionName != null)
+                return false;
+        }
+        else if (!tagDefinitionName.equals(other.tagDefinitionName))
             return false;
-
         return true;
     }
 
-    @Override
-    public int hashCode() {
-        return tagDefinitionName != null ? tagDefinitionName.hashCode() : 0;
-    }
 }