Mercurial > hg > release > thermostat-0.5

changeset 900:916ea399eb6b

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Consolidate Storage.putPojo() to replace and add operations. Reviewed-by: vanaltj Review-thread: http://icedtea.classpath.org/pipermail/thermostat/2012-December/004896.html
author Roman Kennke <rkennke@redhat.com>
date Thu, 10 Jan 2013 14:53:32 +0100
parents 40be72a1345c
children eef06eccb01b
files client/cli/src/main/java/com/redhat/thermostat/client/cli/internal/VMListFormatter.java common/core/src/main/java/com/redhat/thermostat/common/dao/AgentInfoDAOImpl.java common/core/src/main/java/com/redhat/thermostat/common/dao/BackendInfoDAOImpl.java common/core/src/main/java/com/redhat/thermostat/common/dao/HostInfoDAOImpl.java common/core/src/main/java/com/redhat/thermostat/common/dao/NetworkInterfaceInfoDAOImpl.java common/core/src/main/java/com/redhat/thermostat/common/dao/VmInfoDAOImpl.java common/core/src/test/java/com/redhat/thermostat/common/dao/AgentInfoDAOTest.java common/core/src/test/java/com/redhat/thermostat/common/dao/BackendInfoDAOTest.java common/core/src/test/java/com/redhat/thermostat/common/dao/HostInfoDAOTest.java common/core/src/test/java/com/redhat/thermostat/common/dao/NetworkInterfaceInfoDAOTest.java common/core/src/test/java/com/redhat/thermostat/common/dao/VmInfoDAOTest.java host-cpu/common/src/main/java/com/redhat/thermostat/host/cpu/common/internal/CpuStatDAOImpl.java host-cpu/common/src/test/java/com/redhat/thermostat/host/cpu/common/internal/CpuStatDAOTest.java host-memory/common/src/main/java/com/redhat/thermostat/host/memory/common/internal/MemoryStatDAOImpl.java host-memory/common/src/test/java/com/redhat/thermostat/host/memory/common/internal/MemoryStatDAOTest.java storage/core/src/main/java/com/redhat/thermostat/storage/core/Add.java storage/core/src/main/java/com/redhat/thermostat/storage/core/BasePut.java storage/core/src/main/java/com/redhat/thermostat/storage/core/Put.java storage/core/src/main/java/com/redhat/thermostat/storage/core/QueuedStorage.java storage/core/src/main/java/com/redhat/thermostat/storage/core/Replace.java storage/core/src/main/java/com/redhat/thermostat/storage/core/Storage.java storage/core/src/main/java/com/redhat/thermostat/storage/model/AgentIdPojo.java storage/core/src/main/java/com/redhat/thermostat/storage/model/BasePojo.java storage/core/src/test/java/com/redhat/thermostat/storage/core/BasePutTest.java storage/core/src/test/java/com/redhat/thermostat/storage/core/QueuedStorageTest.java storage/mongo/src/main/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorage.java storage/mongo/src/test/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorageTest.java thread/collector/src/main/java/com/redhat/thermostat/thread/dao/impl/ThreadDaoImpl.java thread/collector/src/test/java/com/redhat/thermostat/thread/dao/impl/ThreadDaoImplTest.java vm-classstat/common/src/main/java/com/redhat/thermostat/vm/classstat/common/internal/VmClassStatDAOImpl.java vm-classstat/common/src/test/java/com/redhat/thermostat/vm/classstat/common/internal/VmClassStatDAOTest.java vm-cpu/common/src/main/java/com/redhat/thermostat/vm/cpu/common/internal/VmCpuStatDAOImpl.java vm-cpu/common/src/test/java/com/redhat/thermostat/vm/cpu/common/internal/VmCpuStatDAOTest.java vm-gc/common/src/main/java/com/redhat/thermostat/vm/gc/common/internal/VmGcStatDAOImpl.java vm-gc/common/src/test/java/com/redhat/thermostat/vm/gc/common/internal/VmGcStatDAOTest.java vm-heap-analysis/common/src/main/java/com/redhat/thermostat/vm/heap/analysis/common/internal/HeapDAOImpl.java vm-heap-analysis/common/src/test/java/com/redhat/thermostat/vm/heap/analysis/common/internal/HeapDAOTest.java vm-memory/common/src/main/java/com/redhat/thermostat/vm/memory/common/internal/VmMemoryStatDAOImpl.java vm-memory/common/src/test/java/com/redhat/thermostat/vm/memory/common/internal/VmMemoryStatDAOTest.java web/client/src/main/java/com/redhat/thermostat/web/client/internal/WebStorage.java web/client/src/test/java/com/redhat/thermostat/web/client/internal/WebStorageTest.java web/common/src/main/java/com/redhat/thermostat/web/common/WebInsert.java web/server/src/main/java/com/redhat/thermostat/web/server/WebStorageEndPoint.java web/server/src/test/java/com/redhat/thermostat/web/server/WebStorageEndpointTest.java
diffstat 44 files changed, 702 insertions(+), 166 deletions(-) [+]
line wrap: on
line diff
--- a/client/cli/src/main/java/com/redhat/thermostat/client/cli/internal/VMListFormatter.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/client/cli/src/main/java/com/redhat/thermostat/client/cli/internal/VMListFormatter.java	Thu Jan 10 14:53:32 2013 +0100
@@ -76,6 +76,7 @@
     }
 
     private void printVM(VmRef vm, VmInfo info) {
+        System.err.println("agentId: " + vm.getAgent().getAgentId());
         printLine(vm.getAgent().getAgentId(),
                   vm.getAgent().getHostName(),
                   vm.getId().toString(),
--- a/common/core/src/main/java/com/redhat/thermostat/common/dao/AgentInfoDAOImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/common/core/src/main/java/com/redhat/thermostat/common/dao/AgentInfoDAOImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -41,6 +41,7 @@
 
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Query;
 import com.redhat.thermostat.storage.core.Remove;
 import com.redhat.thermostat.storage.core.Storage;
@@ -104,7 +105,9 @@
 
     @Override
     public void addAgentInformation(AgentInformation agentInfo) {
-        storage.putPojo(AgentInfoDAO.CATEGORY, true, agentInfo);
+        Put replace = storage.createReplace(CATEGORY);
+        replace.setPojo(agentInfo);
+        replace.apply();
     }
 
     @Override
--- a/common/core/src/main/java/com/redhat/thermostat/common/dao/BackendInfoDAOImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/common/core/src/main/java/com/redhat/thermostat/common/dao/BackendInfoDAOImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -43,6 +43,7 @@
 
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Query;
 import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Remove;
@@ -93,7 +94,9 @@
 
     @Override
     public void addBackendInformation(BackendInformation info) {
-        storage.putPojo(BackendInfoDAO.CATEGORY, false, info);
+        Put add = storage.createAdd(BackendInfoDAO.CATEGORY);
+        add.setPojo(info);
+        add.apply();
     }
 
     @Override
--- a/common/core/src/main/java/com/redhat/thermostat/common/dao/HostInfoDAOImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/common/core/src/main/java/com/redhat/thermostat/common/dao/HostInfoDAOImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -42,6 +42,7 @@
 
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Query;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.core.Query.Criteria;
@@ -71,7 +72,9 @@
 
     @Override
     public void putHostInfo(HostInfo info) {
-        storage.putPojo(hostInfoCategory, false, info);
+        Put add = storage.createAdd(hostInfoCategory);
+        add.setPojo(info);
+        add.apply();
     }
 
     @Override
--- a/common/core/src/main/java/com/redhat/thermostat/common/dao/NetworkInterfaceInfoDAOImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/common/core/src/main/java/com/redhat/thermostat/common/dao/NetworkInterfaceInfoDAOImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -41,9 +41,10 @@
 
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Query;
+import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Storage;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.model.NetworkInterfaceInfo;
 
 class NetworkInterfaceInfoDAOImpl implements NetworkInterfaceInfoDAO {
@@ -72,7 +73,9 @@
 
     @Override
     public void putNetworkInterfaceInfo(NetworkInterfaceInfo info) {
-        storage.putPojo(networkInfoCategory, true, info);
+        Put replace = storage.createReplace(networkInfoCategory);
+        replace.setPojo(info);
+        replace.apply();
     }
 
 }
--- a/common/core/src/main/java/com/redhat/thermostat/common/dao/VmInfoDAOImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/common/core/src/main/java/com/redhat/thermostat/common/dao/VmInfoDAOImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -42,10 +42,11 @@
 
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Query;
+import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.core.Update;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.model.VmInfo;
 
 class VmInfoDAOImpl implements VmInfoDAO {
@@ -111,7 +112,9 @@
 
     @Override
     public void putVmInfo(VmInfo info) {
-        storage.putPojo(vmInfoCategory, true, info);
+        Put replace = storage.createReplace(vmInfoCategory);
+        replace.setPojo(info);
+        replace.apply();
     }
 
     @Override
--- a/common/core/src/test/java/com/redhat/thermostat/common/dao/AgentInfoDAOTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/common/core/src/test/java/com/redhat/thermostat/common/dao/AgentInfoDAOTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -56,10 +56,11 @@
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
 import com.redhat.thermostat.storage.core.Query;
+import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Remove;
+import com.redhat.thermostat.storage.core.Replace;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.core.Update;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.model.AgentInformation;
 import com.redhat.thermostat.test.MockQuery;
 
@@ -196,12 +197,16 @@
     @Test
     public void verifyAddAgentInformation() {
         Storage storage = mock(Storage.class);
+        Replace replace = mock(Replace.class);
+        when(storage.createReplace(any(Category.class))).thenReturn(replace);
+
         AgentInfoDAO dao = new AgentInfoDAOImpl(storage);
 
         dao.addAgentInformation(agentInfo1);
 
-        verify(storage).putPojo(AgentInfoDAO.CATEGORY, true, agentInfo1);
-
+        verify(storage).createReplace(AgentInfoDAO.CATEGORY);
+        verify(replace).setPojo(agentInfo1);
+        verify(replace).apply();
     }
 
     @Test
--- a/common/core/src/test/java/com/redhat/thermostat/common/dao/BackendInfoDAOTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/common/core/src/test/java/com/redhat/thermostat/common/dao/BackendInfoDAOTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -38,6 +38,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -51,12 +52,13 @@
 import org.junit.Test;
 import org.mockito.InOrder;
 
+import com.redhat.thermostat.storage.core.Add;
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Remove;
 import com.redhat.thermostat.storage.core.Storage;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.model.BackendInformation;
 import com.redhat.thermostat.test.MockQuery;
 
@@ -109,11 +111,16 @@
     @Test
     public void verifyAddBackendInformation() {
         Storage storage = mock(Storage.class);
+        Add add = mock(Add.class);
+        when(storage.createAdd(any(Category.class))).thenReturn(add);
+
         BackendInfoDAO dao = new BackendInfoDAOImpl(storage);
 
         dao.addBackendInformation(backendInfo1);
 
-        verify(storage).putPojo(BackendInfoDAO.CATEGORY, false, backendInfo1);
+        verify(storage).createAdd(BackendInfoDAO.CATEGORY);
+        verify(add).setPojo(backendInfo1);
+        verify(add).apply();
     }
 
     @Test
--- a/common/core/src/test/java/com/redhat/thermostat/common/dao/HostInfoDAOTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/common/core/src/test/java/com/redhat/thermostat/common/dao/HostInfoDAOTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -53,6 +53,7 @@
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
 
+import com.redhat.thermostat.storage.core.Add;
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
@@ -188,13 +189,18 @@
     @Test
     public void testPutHostInfo() {
         Storage storage = mock(Storage.class);
+        Add add = mock(Add.class);
+        when(storage.createAdd(any(Category.class))).thenReturn(add);
+
         AgentInfoDAO agentInfo = mock(AgentInfoDAO.class);
 
         HostInfo info = new HostInfo(HOST_NAME, OS_NAME, OS_KERNEL, CPU_MODEL, CPU_NUM, MEMORY_TOTAL);
         HostInfoDAO dao = new HostInfoDAOImpl(storage, agentInfo);
         dao.putHostInfo(info);
 
-        verify(storage).putPojo(HostInfoDAO.hostInfoCategory, false, info);
+        verify(storage).createAdd(HostInfoDAO.hostInfoCategory);
+        verify(add).setPojo(info);
+        verify(add).apply();
     }
 
     @Test
--- a/common/core/src/test/java/com/redhat/thermostat/common/dao/NetworkInterfaceInfoDAOTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/common/core/src/test/java/com/redhat/thermostat/common/dao/NetworkInterfaceInfoDAOTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -39,6 +39,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -48,8 +49,10 @@
 
 import org.junit.Test;
 
+import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Replace;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.model.NetworkInterfaceInfo;
 import com.redhat.thermostat.test.MockQuery;
@@ -110,12 +113,17 @@
     @Test
     public void testPutNetworkInterfaceInfo() {
         Storage storage = mock(Storage.class);
+        Replace replace = mock(Replace.class);
+        when(storage.createReplace(any(Category.class))).thenReturn(replace);
+
         NetworkInterfaceInfo info = new NetworkInterfaceInfo(INTERFACE_NAME);
         info.setIp4Addr(IPV4_ADDR);
         info.setIp6Addr(IPV6_ADDR);
         NetworkInterfaceInfoDAO dao = new NetworkInterfaceInfoDAOImpl(storage);
         dao.putNetworkInterfaceInfo(info);
 
-        verify(storage).putPojo(NetworkInterfaceInfoDAO.networkInfoCategory, true, info);
+        verify(storage).createReplace(NetworkInterfaceInfoDAO.networkInfoCategory);
+        verify(replace).setPojo(info);
+        verify(replace).apply();
     }
 }
--- a/common/core/src/test/java/com/redhat/thermostat/common/dao/VmInfoDAOTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/common/core/src/test/java/com/redhat/thermostat/common/dao/VmInfoDAOTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -57,9 +57,10 @@
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
 import com.redhat.thermostat.storage.core.Query;
+import com.redhat.thermostat.storage.core.Query.Criteria;
+import com.redhat.thermostat.storage.core.Replace;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.core.Update;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.model.VmInfo;
 import com.redhat.thermostat.test.MockQuery;
 
@@ -253,13 +254,18 @@
     public void testPutVmInfo() {
 
         Storage storage = mock(Storage.class);
+        Replace replace = mock(Replace.class);
+        when(storage.createReplace(any(Category.class))).thenReturn(replace);
+
         VmInfo info = new VmInfo(vmId, startTime, stopTime, jVersion, jHome,
                 mainClass, commandLine, vmName, vmInfo, vmVersion, vmArgs,
                 props, env, libs);
         VmInfoDAO dao = new VmInfoDAOImpl(storage);
         dao.putVmInfo(info);
 
-        verify(storage).putPojo(VmInfoDAO.vmInfoCategory, true, info);
+        verify(storage).createReplace(VmInfoDAO.vmInfoCategory);
+        verify(replace).setPojo(info);
+        verify(replace).apply();
     }
 
     @Test
--- a/host-cpu/common/src/main/java/com/redhat/thermostat/host/cpu/common/internal/CpuStatDAOImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/host-cpu/common/src/main/java/com/redhat/thermostat/host/cpu/common/internal/CpuStatDAOImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -41,6 +41,7 @@
 import com.redhat.thermostat.common.dao.HostLatestPojoListGetter;
 import com.redhat.thermostat.common.dao.HostRef;
 import com.redhat.thermostat.host.cpu.common.CpuStatDAO;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.model.CpuStat;
 
@@ -63,7 +64,9 @@
 
     @Override
     public void putCpuStat(CpuStat stat) {
-        storage.putPojo(cpuStatCategory, false, stat);
+        Put add = storage.createAdd(cpuStatCategory);
+        add.setPojo(stat);
+        add.apply();
     }
 
     @Override
--- a/host-cpu/common/src/test/java/com/redhat/thermostat/host/cpu/common/internal/CpuStatDAOTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/host-cpu/common/src/test/java/com/redhat/thermostat/host/cpu/common/internal/CpuStatDAOTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -53,13 +53,13 @@
 
 import com.redhat.thermostat.common.dao.HostRef;
 import com.redhat.thermostat.host.cpu.common.CpuStatDAO;
-import com.redhat.thermostat.host.cpu.common.internal.CpuStatDAOImpl;
+import com.redhat.thermostat.storage.core.Add;
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
 import com.redhat.thermostat.storage.core.Query;
+import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Storage;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.model.CpuStat;
 import com.redhat.thermostat.test.MockQuery;
 
@@ -138,11 +138,16 @@
     @Test
     public void testPutCpuStat() {
         Storage storage = mock(Storage.class);
+        Add add = mock(Add.class);
+        when(storage.createAdd(any(Category.class))).thenReturn(add);
+        
         CpuStat stat = new CpuStat(1,  new double[] {5.0, 10.0, 15.0});
         CpuStatDAO dao = new CpuStatDAOImpl(storage);
         dao.putCpuStat(stat);
 
-        verify(storage).putPojo(CpuStatDAO.cpuStatCategory, false, stat);
+        verify(storage).createAdd(CpuStatDAO.cpuStatCategory);
+        verify(add).setPojo(stat);
+        verify(add).apply();
     }
 
     @Test
--- a/host-memory/common/src/main/java/com/redhat/thermostat/host/memory/common/internal/MemoryStatDAOImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/host-memory/common/src/main/java/com/redhat/thermostat/host/memory/common/internal/MemoryStatDAOImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -41,6 +41,7 @@
 import com.redhat.thermostat.common.dao.HostLatestPojoListGetter;
 import com.redhat.thermostat.common.dao.HostRef;
 import com.redhat.thermostat.host.memory.common.MemoryStatDAO;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.model.MemoryStat;
 
@@ -63,7 +64,9 @@
 
     @Override
     public void putMemoryStat(MemoryStat stat) {
-        storage.putPojo(memoryStatCategory, false, stat);
+        Put add = storage.createAdd(memoryStatCategory);
+        add.setPojo(stat);
+        add.apply();
     }
 
     @Override
--- a/host-memory/common/src/test/java/com/redhat/thermostat/host/memory/common/internal/MemoryStatDAOTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/host-memory/common/src/test/java/com/redhat/thermostat/host/memory/common/internal/MemoryStatDAOTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -54,6 +54,7 @@
 
 import com.redhat.thermostat.common.dao.HostRef;
 import com.redhat.thermostat.host.memory.common.MemoryStatDAO;
+import com.redhat.thermostat.storage.core.Add;
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
@@ -135,11 +136,16 @@
     @Test
     public void testPutMemoryStat() {
         Storage storage = mock(Storage.class);
+        Add add = mock(Add.class);
+        when(storage.createAdd(any(Category.class))).thenReturn(add);
+
         MemoryStat stat = new MemoryStat(TIMESTAMP, TOTAL, FREE, BUFFERS, CACHED, SWAP_TOTAL, SWAP_FREE, COMMIT_LIMIT);
         MemoryStatDAO dao = new MemoryStatDAOImpl(storage);
         dao.putMemoryStat(stat);
 
-        verify(storage).putPojo(MemoryStatDAO.memoryStatCategory, false, stat);
+        verify(storage).createAdd(MemoryStatDAO.memoryStatCategory);
+        verify(add).setPojo(stat);
+        verify(add).apply();
     }
 
     @Test
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/core/Add.java	Thu Jan 10 14:53:32 2013 +0100
@@ -0,0 +1,42 @@
+/*
+ * Copyright 2012 Red Hat, Inc.
+ *
+ * This file is part of Thermostat.
+ *
+ * Thermostat is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published
+ * by the Free Software Foundation; either version 2, or (at your
+ * option) any later version.
+ *
+ * Thermostat is distributed in the hope that it will be useful, but
+ * WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Thermostat; see the file COPYING.  If not see
+ * <http://www.gnu.org/licenses/>.
+ *
+ * Linking this code with other modules is making a combined work
+ * based on this code.  Thus, the terms and conditions of the GNU
+ * General Public License cover the whole combination.
+ *
+ * As a special exception, the copyright holders of this code give
+ * you permission to link this code with independent modules to
+ * produce an executable, regardless of the license terms of these
+ * independent modules, and to copy and distribute the resulting
+ * executable under terms of your choice, provided that you also
+ * meet, for each linked independent module, the terms and conditions
+ * of the license of that module.  An independent module is a module
+ * which is not derived from or based on this code.  If you modify
+ * this code, you may extend this exception to your version of the
+ * library, but you are not obligated to do so.  If you do not wish
+ * to do so, delete this exception statement from your version.
+ */
+
+
+package com.redhat.thermostat.storage.core;
+
+public interface Add extends Put {
+
+}
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/core/BasePut.java	Thu Jan 10 14:53:32 2013 +0100
@@ -0,0 +1,68 @@
+/*
+ * Copyright 2012 Red Hat, Inc.
+ *
+ * This file is part of Thermostat.
+ *
+ * Thermostat is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published
+ * by the Free Software Foundation; either version 2, or (at your
+ * option) any later version.
+ *
+ * Thermostat is distributed in the hope that it will be useful, but
+ * WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Thermostat; see the file COPYING.  If not see
+ * <http://www.gnu.org/licenses/>.
+ *
+ * Linking this code with other modules is making a combined work
+ * based on this code.  Thus, the terms and conditions of the GNU
+ * General Public License cover the whole combination.
+ *
+ * As a special exception, the copyright holders of this code give
+ * you permission to link this code with independent modules to
+ * produce an executable, regardless of the license terms of these
+ * independent modules, and to copy and distribute the resulting
+ * executable under terms of your choice, provided that you also
+ * meet, for each linked independent module, the terms and conditions
+ * of the license of that module.  An independent module is a module
+ * which is not derived from or based on this code.  If you modify
+ * this code, you may extend this exception to your version of the
+ * library, but you are not obligated to do so.  If you do not wish
+ * to do so, delete this exception statement from your version.
+ */
+
+
+package com.redhat.thermostat.storage.core;
+
+import com.redhat.thermostat.storage.model.Pojo;
+
+
+public abstract class BasePut implements Put {
+
+    private Category category;
+    private Pojo pojo;
+
+    public final void setCategory(Category category) {
+        if (this.category != null) {
+            throw new IllegalStateException();
+        }
+        this.category = category;
+    }
+
+    public final Category getCategory() {
+        return category;
+    }
+
+    @Override
+    public final void setPojo(Pojo pojo) {
+        this.pojo = pojo;
+    }
+
+    public final Pojo getPojo() {
+        return pojo;
+    }
+
+}
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/core/Put.java	Thu Jan 10 14:53:32 2013 +0100
@@ -0,0 +1,59 @@
+/*
+ * Copyright 2012 Red Hat, Inc.
+ *
+ * This file is part of Thermostat.
+ *
+ * Thermostat is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published
+ * by the Free Software Foundation; either version 2, or (at your
+ * option) any later version.
+ *
+ * Thermostat is distributed in the hope that it will be useful, but
+ * WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Thermostat; see the file COPYING.  If not see
+ * <http://www.gnu.org/licenses/>.
+ *
+ * Linking this code with other modules is making a combined work
+ * based on this code.  Thus, the terms and conditions of the GNU
+ * General Public License cover the whole combination.
+ *
+ * As a special exception, the copyright holders of this code give
+ * you permission to link this code with independent modules to
+ * produce an executable, regardless of the license terms of these
+ * independent modules, and to copy and distribute the resulting
+ * executable under terms of your choice, provided that you also
+ * meet, for each linked independent module, the terms and conditions
+ * of the license of that module.  An independent module is a module
+ * which is not derived from or based on this code.  If you modify
+ * this code, you may extend this exception to your version of the
+ * library, but you are not obligated to do so.  If you do not wish
+ * to do so, delete this exception statement from your version.
+ */
+
+
+package com.redhat.thermostat.storage.core;
+
+import com.redhat.thermostat.storage.model.Pojo;
+
+
+public interface Put {
+
+    /**
+     * Sets the POJO that is to be put into the database.
+     * Notice that this is usually not necessary to call, because Storage.add() 
+     * and Storage.replace() already take this argument. It is useful
+     * for pre-building and caching queries though.
+     *
+     * @param the pojo to be put into the database
+     */
+    void setPojo(Pojo pojo);
+
+    /**
+     * Applies this put operation to the database.
+     */
+    void apply();
+}
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/core/QueuedStorage.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/core/QueuedStorage.java	Thu Jan 10 14:53:32 2013 +0100
@@ -42,11 +42,28 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
-import com.redhat.thermostat.storage.model.AgentIdPojo;
 import com.redhat.thermostat.storage.model.Pojo;
 
 public final class QueuedStorage implements Storage {
 
+    private class QueuedReplace extends BasePut implements Replace {
+
+        @Override
+        public void apply() {
+            replaceImpl(getCategory(), getPojo());
+        }
+        
+    }
+
+    private class QueuedAdd extends BasePut implements Add {
+
+        @Override
+        public void apply() {
+            addImpl(getCategory(), getPojo());
+        }
+        
+    }
+
     private Storage delegate;
     private ExecutorService executor;
     private ExecutorService fileExecutor;
@@ -79,13 +96,43 @@
     }
 
     @Override
-    public void putPojo(final Category category, final boolean replace, final AgentIdPojo pojo) {
+    public Add createAdd(Category into) {
+        QueuedAdd add = new QueuedAdd();
+        add.setCategory(into);
+        return add;
+    }
 
+    @Override
+    public Replace createReplace(Category into) {
+        QueuedReplace replace = new QueuedReplace();
+        replace.setCategory(into);
+        return replace;
+    }
+
+    private void replaceImpl(final Category category, final Pojo pojo) {
+        
         executor.execute(new Runnable() {
             
             @Override
             public void run() {
-                delegate.putPojo(category, replace, pojo);
+                Replace replace = delegate.createReplace(category);
+                replace.setPojo(pojo);
+                replace.apply();
+            }
+
+        });
+
+    }
+
+    private void addImpl(final Category category, final Pojo pojo) {
+        
+        executor.execute(new Runnable() {
+            
+            @Override
+            public void run() {
+                Add add = delegate.createAdd(category);
+                add.setPojo(pojo);
+                add.apply();
             }
 
         });
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/core/Replace.java	Thu Jan 10 14:53:32 2013 +0100
@@ -0,0 +1,42 @@
+/*
+ * Copyright 2012 Red Hat, Inc.
+ *
+ * This file is part of Thermostat.
+ *
+ * Thermostat is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published
+ * by the Free Software Foundation; either version 2, or (at your
+ * option) any later version.
+ *
+ * Thermostat is distributed in the hope that it will be useful, but
+ * WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Thermostat; see the file COPYING.  If not see
+ * <http://www.gnu.org/licenses/>.
+ *
+ * Linking this code with other modules is making a combined work
+ * based on this code.  Thus, the terms and conditions of the GNU
+ * General Public License cover the whole combination.
+ *
+ * As a special exception, the copyright holders of this code give
+ * you permission to link this code with independent modules to
+ * produce an executable, regardless of the license terms of these
+ * independent modules, and to copy and distribute the resulting
+ * executable under terms of your choice, provided that you also
+ * meet, for each linked independent module, the terms and conditions
+ * of the license of that module.  An independent module is a module
+ * which is not derived from or based on this code.  If you modify
+ * this code, you may extend this exception to your version of the
+ * library, but you are not obligated to do so.  If you do not wish
+ * to do so, delete this exception statement from your version.
+ */
+
+
+package com.redhat.thermostat.storage.core;
+
+public interface Replace extends Put {
+
+}
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/core/Storage.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/core/Storage.java	Thu Jan 10 14:53:32 2013 +0100
@@ -39,7 +39,6 @@
 import java.io.InputStream;
 import java.util.UUID;
 
-import com.redhat.thermostat.storage.model.AgentIdPojo;
 import com.redhat.thermostat.storage.model.Pojo;
 
 /**
@@ -59,7 +58,8 @@
 
     // TODO why does putPojo have an agent id param but not updatePojo and removePojo?
 
-    void putPojo(Category category, boolean replace, AgentIdPojo pojo);
+    Add createAdd(Category category);
+    Replace createReplace(Category category);
 
     void updatePojo(Update update);
 
@@ -84,6 +84,7 @@
     Update createUpdate();
     Remove createRemove();
 
+
     void shutdown();
 
 }
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/model/AgentIdPojo.java	Fri Dec 21 14:19:06 2012 +0100
+++ /dev/null	Thu Jan 01 00:00:00 1970 +0000
@@ -1,44 +0,0 @@
-/*
- * Copyright 2012 Red Hat, Inc.
- *
- * This file is part of Thermostat.
- *
- * Thermostat is free software; you can redistribute it and/or modify
- * it under the terms of the GNU General Public License as published
- * by the Free Software Foundation; either version 2, or (at your
- * option) any later version.
- *
- * Thermostat is distributed in the hope that it will be useful, but
- * WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
- * General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with Thermostat; see the file COPYING.  If not see
- * <http://www.gnu.org/licenses/>.
- *
- * Linking this code with other modules is making a combined work
- * based on this code.  Thus, the terms and conditions of the GNU
- * General Public License cover the whole combination.
- *
- * As a special exception, the copyright holders of this code give
- * you permission to link this code with independent modules to
- * produce an executable, regardless of the license terms of these
- * independent modules, and to copy and distribute the resulting
- * executable under terms of your choice, provided that you also
- * meet, for each linked independent module, the terms and conditions
- * of the license of that module.  An independent module is a module
- * which is not derived from or based on this code.  If you modify
- * this code, you may extend this exception to your version of the
- * library, but you are not obligated to do so.  If you do not wish
- * to do so, delete this exception statement from your version.
- */
-
-
-package com.redhat.thermostat.storage.model;
-
-public interface AgentIdPojo extends Pojo {
-
-    void setAgentId(String agentId);
-    String getAgentId();
-}
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/model/BasePojo.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/model/BasePojo.java	Thu Jan 10 14:53:32 2013 +0100
@@ -41,7 +41,7 @@
 
 import com.redhat.thermostat.storage.core.Persist;
 
-public class BasePojo implements AgentIdPojo {
+public class BasePojo implements Pojo {
 
     private String agentId;
 
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/test/java/com/redhat/thermostat/storage/core/BasePutTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -0,0 +1,94 @@
+/*
+ * Copyright 2012 Red Hat, Inc.
+ *
+ * This file is part of Thermostat.
+ *
+ * Thermostat is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published
+ * by the Free Software Foundation; either version 2, or (at your
+ * option) any later version.
+ *
+ * Thermostat is distributed in the hope that it will be useful, but
+ * WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Thermostat; see the file COPYING.  If not see
+ * <http://www.gnu.org/licenses/>.
+ *
+ * Linking this code with other modules is making a combined work
+ * based on this code.  Thus, the terms and conditions of the GNU
+ * General Public License cover the whole combination.
+ *
+ * As a special exception, the copyright holders of this code give
+ * you permission to link this code with independent modules to
+ * produce an executable, regardless of the license terms of these
+ * independent modules, and to copy and distribute the resulting
+ * executable under terms of your choice, provided that you also
+ * meet, for each linked independent module, the terms and conditions
+ * of the license of that module.  An independent module is a module
+ * which is not derived from or based on this code.  If you modify
+ * this code, you may extend this exception to your version of the
+ * library, but you are not obligated to do so.  If you do not wish
+ * to do so, delete this exception statement from your version.
+ */
+
+
+package com.redhat.thermostat.storage.core;
+
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import com.redhat.thermostat.storage.model.Pojo;
+
+public class BasePutTest {
+
+    private BasePut insert;
+
+    @Before
+    public void setUp() {
+        insert = new BasePut() {
+            public void apply() {
+                // Do nothing.
+            }
+        };
+    }
+
+    @After
+    public void tearDown() {
+        insert = null;
+    }
+
+    @Test
+    public void testCategory() {
+        Category category = mock(Category.class);
+
+        assertNull(insert.getCategory());
+        insert.setCategory(category);
+        assertSame(category, insert.getCategory());
+
+        try {
+            insert.setCategory(category);
+            fail();
+        } catch (IllegalStateException ex) {
+            // Ok.
+        }
+    }
+
+    @Test
+    public void testPojo() {
+        Pojo pojo = mock(Pojo.class);
+
+        assertNull(insert.getPojo());
+        insert.setPojo(pojo);
+        assertSame(pojo, insert.getPojo());
+        
+    }
+}
--- a/storage/core/src/test/java/com/redhat/thermostat/storage/core/QueuedStorageTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/storage/core/src/test/java/com/redhat/thermostat/storage/core/QueuedStorageTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -67,7 +67,6 @@
 import org.junit.Before;
 import org.junit.Test;
 
-import com.redhat.thermostat.storage.model.AgentIdPojo;
 import com.redhat.thermostat.storage.model.Pojo;
 
 
@@ -165,6 +164,7 @@
             // Not used.
             return null;
         }
+
     }
 
     private static class TestPojo implements Pojo {
@@ -173,6 +173,8 @@
 
     private QueuedStorage queuedStorage;
     private Storage delegateStorage;
+    private Add delegateAdd;
+    private Replace delegateReplace;
 
     private TestExecutor executor;
     private TestExecutor fileExecutor;
@@ -188,9 +190,15 @@
         executor = new TestExecutor();
         fileExecutor = new TestExecutor();
         delegateStorage = mock(Storage.class);
+
+        delegateAdd = mock(Add.class);
+        delegateReplace = mock(Replace.class);
+
         Update update = mock(Update.class);
         Remove remove = mock(Remove.class);
         Query query = mock(Query.class);
+        when(delegateStorage.createAdd(any(Category.class))).thenReturn(delegateAdd);
+        when(delegateStorage.createReplace(any(Category.class))).thenReturn(delegateReplace);
         when(delegateStorage.createUpdate()).thenReturn(update);
         when(delegateStorage.createRemove()).thenReturn(remove);
         when(delegateStorage.createQuery()).thenReturn(query);
@@ -218,17 +226,23 @@
     }
 
     @Test
-    public void testPutPojo() {
+    public void testInsert() {
         Category category = mock(Category.class);
-        AgentIdPojo pojo = mock(AgentIdPojo.class);
+        Pojo pojo = mock(Pojo.class);
 
-        queuedStorage.putPojo(category, true, pojo);
+        Put put = queuedStorage.createReplace(category);
+        put.setPojo(pojo);
+        put.apply();
 
         Runnable r = executor.getTask();
         assertNotNull(r);
         verifyZeroInteractions(delegateStorage);
+        verifyZeroInteractions(delegateReplace);
+
         r.run();
-        verify(delegateStorage, times(1)).putPojo(category, true, pojo);
+        verify(delegateStorage).createReplace(category);
+        verify(delegateReplace).setPojo(pojo);
+        verify(delegateReplace).apply();
         verifyNoMoreInteractions(delegateStorage);
 
         assertNull(fileExecutor.getTask());
--- a/storage/mongo/src/main/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorage.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/storage/mongo/src/main/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorage.java	Thu Jan 10 14:53:32 2013 +0100
@@ -53,6 +53,8 @@
 import com.mongodb.gridfs.GridFSInputFile;
 import com.redhat.thermostat.storage.config.StartupConfiguration;
 import com.redhat.thermostat.storage.core.AbstractQuery.Sort;
+import com.redhat.thermostat.storage.core.Add;
+import com.redhat.thermostat.storage.core.BasePut;
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Connection;
 import com.redhat.thermostat.storage.core.Connection.ConnectionListener;
@@ -61,9 +63,9 @@
 import com.redhat.thermostat.storage.core.Key;
 import com.redhat.thermostat.storage.core.Query;
 import com.redhat.thermostat.storage.core.Remove;
+import com.redhat.thermostat.storage.core.Replace;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.core.Update;
-import com.redhat.thermostat.storage.model.AgentIdPojo;
 import com.redhat.thermostat.storage.model.Pojo;
 
 /**
@@ -75,6 +77,24 @@
 
     public static final String SET_MODIFIER = "$set";
 
+    private class MongoAdd extends BasePut implements Add {
+
+        @Override
+        public void apply() {
+            addImpl(getCategory(), getPojo());
+        }
+        
+    }
+
+    private class MongoReplace extends BasePut implements Replace {
+
+        @Override
+        public void apply() {
+            replaceImpl(getCategory(), getPojo());
+        }
+        
+    }
+
     private MongoConnection conn;
     private DB db = null;
     private Map<String, DBCollection> collectionCache = new HashMap<String, DBCollection>();
@@ -113,45 +133,48 @@
         return agentId.toString();
     }
 
-    private String getAgentQueryKeyFromGlobalAgent() {
-        if (agentId != null) {
-            return agentId.toString();
-        } else {
-            return null;
-        }
-    }
-
-    private String getAgentQueryKeyFromChunkOrGlobalAgent(AgentIdPojo pojo) {
-        String queryKey = getAgentQueryKeyFromGlobalAgent();
-        if (queryKey != null) {
-            return queryKey;
-        } else {
-            return pojo.getAgentId();
-        }
+    @Override
+    public Add createAdd(Category into) {
+        MongoAdd add = new MongoAdd();
+        add.setCategory(into);
+        return add;
     }
 
     @Override
-    public void putPojo(final Category cat, final boolean replace, final AgentIdPojo pojo) {
+    public Replace createReplace(Category into) {
+        MongoReplace replace = new MongoReplace();
+        replace.setCategory(into);
+        return replace;
+    }
+
+    private void addImpl(final Category cat, final Pojo pojo) {
+        DBCollection coll = getCachedCollection(cat);
+        DBObject toInsert = preparePut(pojo);
+        coll.insert(toInsert);
+    }
+
+    private void replaceImpl(final Category cat, final Pojo pojo) {
         DBCollection coll = getCachedCollection(cat);
+        DBObject toInsert = preparePut(pojo);
+
+        DBObject query = new BasicDBObject();
+        Collection<Key<?>> keys = cat.getKeys();
+        for (Key<?> key : keys) {
+            if (key.isPartialCategoryKey()) {
+                String name = key.getName();
+                query.put(name, toInsert.get(name));
+            }
+        }
+        coll.update(query, toInsert, true, false);
+    }
+
+    private DBObject preparePut(final Pojo pojo) {
         MongoPojoConverter converter = new MongoPojoConverter();
         DBObject toInsert = converter.convertPojoToMongo(pojo);
-        String agentId = getAgentQueryKeyFromChunkOrGlobalAgent(pojo);
-        toInsert.put(Key.AGENT_ID.getName(), agentId);
-        if (replace) {
-            // TODO: Split this part out into a separate method. It is a very bad practice to
-            // completely change the behaviour of a method based on a boolean flag.
-            DBObject query = new BasicDBObject();
-            Collection<Key<?>> keys = cat.getKeys();
-            for (Key<?> key : keys) {
-                if (key.isPartialCategoryKey()) {
-                    String name = key.getName();
-                    query.put(name, toInsert.get(name));
-                }
-            }
-            coll.update(query, toInsert, true, false);
-        } else {
-            coll.insert(toInsert);
+        if (toInsert.get(Key.AGENT_ID.getName()) == null) {
+            toInsert.put(Key.AGENT_ID.getName(), getAgentId());
         }
+        return toInsert;
     }
 
     @Override
@@ -173,11 +196,6 @@
         Category category = mongoRemove.getCategory();
         DBCollection coll = getCachedCollection(category);
 
-        String agentId = getAgentQueryKeyFromGlobalAgent();
-        if (agentId != null) {
-            query.put(Key.AGENT_ID.getName(), agentId);
-        }
-
         coll.remove(query);
     }
 
@@ -200,8 +218,7 @@
 
     @Override
     public void purge() {
-        String deleteKey = getAgentQueryKeyFromGlobalAgent();
-        BasicDBObject query = new BasicDBObject(Key.AGENT_ID.getName(), deleteKey);
+        BasicDBObject query = new BasicDBObject(Key.AGENT_ID.getName(), getAgentId());
         for (DBCollection coll : collectionCache.values()) {
             coll.remove(query);
         }
--- a/storage/mongo/src/test/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorageTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/storage/mongo/src/test/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorageTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -74,11 +74,13 @@
 import com.mongodb.gridfs.GridFSDBFile;
 import com.mongodb.gridfs.GridFSInputFile;
 import com.redhat.thermostat.storage.config.StartupConfiguration;
+import com.redhat.thermostat.storage.core.Add;
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Entity;
 import com.redhat.thermostat.storage.core.Key;
 import com.redhat.thermostat.storage.core.Persist;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Query;
 import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Update;
@@ -413,7 +415,9 @@
         MongoStorage storage = makeStorage();
         TestClass pojo = new TestClass();
         pojo.setAgentId("123");
-        storage.putPojo(testCategory, false, pojo);
+        Put add = storage.createAdd(testCategory);
+        add.setPojo(pojo);
+        add.apply();
         ArgumentCaptor<DBObject> dbobj = ArgumentCaptor.forClass(DBObject.class);
         verify(testCollection).insert(dbobj.capture());
         DBObject val = dbobj.getValue();
@@ -426,8 +430,9 @@
         MongoStorage storage = makeStorage();
         storage.setAgentId(new UUID(1, 2));
         TestClass pojo = new TestClass();
-        pojo.setAgentId("123");
-        storage.putPojo(testCategory, false, pojo);
+        Add add = storage.createAdd(testCategory);
+        add.setPojo(pojo);
+        add.apply();
         ArgumentCaptor<DBObject> dbobj = ArgumentCaptor.forClass(DBObject.class);
         verify(testCollection).insert(dbobj.capture());
         DBObject val = dbobj.getValue();
@@ -500,7 +505,9 @@
         pojo.setKey5("test5");
 
         MongoStorage storage = makeStorage();
-        storage.putPojo(testCategory, true, pojo);
+        Put replace = storage.createReplace(testCategory);
+        replace.setPojo(pojo);
+        replace.apply();
 
         ArgumentCaptor<DBObject> queryCaptor = ArgumentCaptor.forClass(DBObject.class);
         ArgumentCaptor<DBObject> valueCaptor = ArgumentCaptor.forClass(DBObject.class);
--- a/thread/collector/src/main/java/com/redhat/thermostat/thread/dao/impl/ThreadDaoImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/thread/collector/src/main/java/com/redhat/thermostat/thread/dao/impl/ThreadDaoImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -43,10 +43,10 @@
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Query;
+import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Storage;
-import com.redhat.thermostat.storage.core.Update;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.thread.dao.ThreadDao;
 import com.redhat.thermostat.thread.model.ThreadInfoData;
 import com.redhat.thermostat.thread.model.ThreadSummary;
@@ -81,12 +81,16 @@
     
     @Override
     public void saveCapabilities(VMThreadCapabilities caps) {
-        storage.putPojo(THREAD_CAPABILITIES, true, caps);
+        Put replace = storage.createReplace(THREAD_CAPABILITIES);
+        replace.setPojo(caps);
+        replace.apply();
     }
     
     @Override
     public void saveSummary(ThreadSummary summary) {
-        storage.putPojo(THREAD_SUMMARY, false, summary);
+        Put add = storage.createAdd(THREAD_SUMMARY);
+        add.setPojo(summary);
+        add.apply();
     }
     
     @Override
@@ -121,7 +125,9 @@
     
     @Override
     public void saveThreadInfo(ThreadInfoData info) {
-        storage.putPojo(THREAD_INFO, false, info);
+        Put add = storage.createAdd(THREAD_INFO);
+        add.setPojo(info);
+        add.apply();
     }
 
     @Override
--- a/thread/collector/src/test/java/com/redhat/thermostat/thread/dao/impl/ThreadDaoImplTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/thread/collector/src/test/java/com/redhat/thermostat/thread/dao/impl/ThreadDaoImplTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -38,6 +38,7 @@
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -46,9 +47,11 @@
 
 import com.redhat.thermostat.common.dao.HostRef;
 import com.redhat.thermostat.common.dao.VmRef;
+import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Query.Criteria;
+import com.redhat.thermostat.storage.core.Replace;
 import com.redhat.thermostat.storage.core.Storage;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.test.MockQuery;
 import com.redhat.thermostat.thread.dao.ThreadDao;
 import com.redhat.thermostat.thread.model.VMThreadCapabilities;
@@ -98,7 +101,9 @@
     @Test
     public void testSaveVMCapabilities() {
         Storage storage = mock(Storage.class);
-        
+        Replace replace = mock(Replace.class);
+        when(storage.createReplace(any(Category.class))).thenReturn(replace);
+
         VMThreadCapabilities caps = mock(VMThreadCapabilities.class);
         when(caps.supportContentionMonitor()).thenReturn(true);
         when(caps.supportCPUTime()).thenReturn(true);
@@ -107,7 +112,9 @@
         ThreadDaoImpl dao = new ThreadDaoImpl(storage);
         dao.saveCapabilities(caps);
 
-        verify(storage).putPojo(ThreadDao.THREAD_CAPABILITIES, true, caps);
+        verify(storage).createReplace(ThreadDao.THREAD_CAPABILITIES);
+        verify(replace).setPojo(caps);
+        verify(replace).apply();
 
     }
 }
--- a/vm-classstat/common/src/main/java/com/redhat/thermostat/vm/classstat/common/internal/VmClassStatDAOImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/vm-classstat/common/src/main/java/com/redhat/thermostat/vm/classstat/common/internal/VmClassStatDAOImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -40,6 +40,7 @@
 
 import com.redhat.thermostat.common.dao.VmLatestPojoListGetter;
 import com.redhat.thermostat.common.dao.VmRef;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.model.VmClassStat;
 import com.redhat.thermostat.vm.classstat.common.VmClassStatDAO;
@@ -62,6 +63,8 @@
 
     @Override
     public void putVmClassStat(VmClassStat stat) {
-        storage.putPojo(vmClassStatsCategory, false, stat);
+        Put add = storage.createAdd(vmClassStatsCategory);
+        add.setPojo(stat);
+        add.apply();
     }
 }
--- a/vm-classstat/common/src/test/java/com/redhat/thermostat/vm/classstat/common/internal/VmClassStatDAOTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/vm-classstat/common/src/test/java/com/redhat/thermostat/vm/classstat/common/internal/VmClassStatDAOTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -54,15 +54,16 @@
 
 import com.redhat.thermostat.common.dao.HostRef;
 import com.redhat.thermostat.common.dao.VmRef;
+import com.redhat.thermostat.storage.core.Add;
+import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
 import com.redhat.thermostat.storage.core.Query;
+import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Storage;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.model.VmClassStat;
 import com.redhat.thermostat.test.MockQuery;
 import com.redhat.thermostat.vm.classstat.common.VmClassStatDAO;
-import com.redhat.thermostat.vm.classstat.common.internal.VmClassStatDAOImpl;
 
 public class VmClassStatDAOTest {
 
@@ -129,10 +130,15 @@
     public void testPutVmClassStat() {
 
         Storage storage = mock(Storage.class);
+        Add add = mock(Add.class);
+        when(storage.createAdd(any(Category.class))).thenReturn(add);
+
         VmClassStat stat = new VmClassStat(VM_ID, TIMESTAMP, LOADED_CLASSES);
         VmClassStatDAO dao = new VmClassStatDAOImpl(storage);
         dao.putVmClassStat(stat);
 
-        verify(storage).putPojo(VmClassStatDAO.vmClassStatsCategory, false, stat);
+        verify(storage).createAdd(VmClassStatDAO.vmClassStatsCategory);
+        verify(add).setPojo(stat);
+        verify(add).apply();
     }
 }
--- a/vm-cpu/common/src/main/java/com/redhat/thermostat/vm/cpu/common/internal/VmCpuStatDAOImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/vm-cpu/common/src/main/java/com/redhat/thermostat/vm/cpu/common/internal/VmCpuStatDAOImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -40,6 +40,7 @@
 
 import com.redhat.thermostat.common.dao.VmLatestPojoListGetter;
 import com.redhat.thermostat.common.dao.VmRef;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.model.VmCpuStat;
 import com.redhat.thermostat.vm.cpu.common.VmCpuStatDAO;
@@ -62,6 +63,8 @@
 
     @Override
     public void putVmCpuStat(VmCpuStat stat) {
-        storage.putPojo(vmCpuStatCategory, false, stat);
+        Put add = storage.createAdd(vmCpuStatCategory);
+        add.setPojo(stat);
+        add.apply();
     }
 }
--- a/vm-cpu/common/src/test/java/com/redhat/thermostat/vm/cpu/common/internal/VmCpuStatDAOTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/vm-cpu/common/src/test/java/com/redhat/thermostat/vm/cpu/common/internal/VmCpuStatDAOTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -55,11 +55,13 @@
 
 import com.redhat.thermostat.common.dao.HostRef;
 import com.redhat.thermostat.common.dao.VmRef;
+import com.redhat.thermostat.storage.core.Add;
+import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
 import com.redhat.thermostat.storage.core.Query;
+import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Storage;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.model.VmCpuStat;
 import com.redhat.thermostat.test.MockQuery;
 import com.redhat.thermostat.vm.cpu.common.VmCpuStatDAO;
@@ -130,11 +132,16 @@
     @Test
     public void testPutVmCpuStat() {
         Storage storage = mock(Storage.class);
+        Add add = mock(Add.class);
+        when(storage.createAdd(any(Category.class))).thenReturn(add);
+        
         VmCpuStat stat = new VmCpuStat(TIMESTAMP, VM_ID, CPU_LOAD);
         VmCpuStatDAO dao = new VmCpuStatDAOImpl(storage);
         dao.putVmCpuStat(stat);
 
-        verify(storage).putPojo(VmCpuStatDAO.vmCpuStatCategory, false, stat);
+        verify(storage).createAdd(VmCpuStatDAO.vmCpuStatCategory);
+        verify(add).setPojo(stat);
+        verify(add).apply();
 
     }
 }
--- a/vm-gc/common/src/main/java/com/redhat/thermostat/vm/gc/common/internal/VmGcStatDAOImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/vm-gc/common/src/main/java/com/redhat/thermostat/vm/gc/common/internal/VmGcStatDAOImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -40,6 +40,7 @@
 
 import com.redhat.thermostat.common.dao.VmLatestPojoListGetter;
 import com.redhat.thermostat.common.dao.VmRef;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.model.VmGcStat;
 import com.redhat.thermostat.vm.gc.common.VmGcStatDAO;
@@ -62,7 +63,9 @@
 
     @Override
     public void putVmGcStat(VmGcStat stat) {
-        storage.putPojo(vmGcStatCategory, false, stat);
+        Put add = storage.createAdd(vmGcStatCategory);
+        add.setPojo(stat);
+        add.apply();
     }
 
 }
--- a/vm-gc/common/src/test/java/com/redhat/thermostat/vm/gc/common/internal/VmGcStatDAOTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/vm-gc/common/src/test/java/com/redhat/thermostat/vm/gc/common/internal/VmGcStatDAOTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -52,15 +52,16 @@
 
 import com.redhat.thermostat.common.dao.HostRef;
 import com.redhat.thermostat.common.dao.VmRef;
+import com.redhat.thermostat.storage.core.Add;
+import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
 import com.redhat.thermostat.storage.core.Query;
+import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Storage;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.model.VmGcStat;
 import com.redhat.thermostat.test.MockQuery;
 import com.redhat.thermostat.vm.gc.common.VmGcStatDAO;
-import com.redhat.thermostat.vm.gc.common.internal.VmGcStatDAOImpl;
 
 public class VmGcStatDAOTest {
 
@@ -124,10 +125,15 @@
     @Test
     public void testPutVmGcStat() {
         Storage storage = mock(Storage.class);
+        Add add = mock(Add.class);
+        when(storage.createAdd(any(Category.class))).thenReturn(add);
+        
         VmGcStat stat = new VmGcStat(VM_ID, TIMESTAMP, COLLECTOR, RUN_COUNT, WALL_TIME);
         VmGcStatDAO dao = new VmGcStatDAOImpl(storage);
         dao.putVmGcStat(stat);
 
-        verify(storage).putPojo(VmGcStatDAO.vmGcStatCategory, false, stat);
+        verify(storage).createAdd(VmGcStatDAO.vmGcStatCategory);
+        verify(add).setPojo(stat);
+        verify(add).apply();
     }
 }
--- a/vm-heap-analysis/common/src/main/java/com/redhat/thermostat/vm/heap/analysis/common/internal/HeapDAOImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/vm-heap-analysis/common/src/main/java/com/redhat/thermostat/vm/heap/analysis/common/internal/HeapDAOImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -53,6 +53,7 @@
 import com.redhat.thermostat.common.utils.LoggingUtils;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Query;
 import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Storage;
@@ -86,7 +87,10 @@
         if (histogramData != null) {
             heapInfo.setHistogramId(histogramId);
         }
-        storage.putPojo(heapInfoCategory, false, heapInfo);
+        Put add = storage.createAdd(heapInfoCategory);
+        add.setPojo(heapInfo);
+        add.apply();
+
         if (heapDumpData != null) {
             storage.saveFile(heapDumpId, new FileInputStream(heapDumpData));
         }
--- a/vm-heap-analysis/common/src/test/java/com/redhat/thermostat/vm/heap/analysis/common/internal/HeapDAOTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/vm-heap-analysis/common/src/test/java/com/redhat/thermostat/vm/heap/analysis/common/internal/HeapDAOTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -66,6 +66,7 @@
 
 import com.redhat.thermostat.common.dao.HostRef;
 import com.redhat.thermostat.common.dao.VmRef;
+import com.redhat.thermostat.storage.core.Add;
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
@@ -85,6 +86,7 @@
 
     private HeapDAO dao;
     private Storage storage;
+    private Add add;
     private HeapInfo heapInfo;
     private File heapDumpData;
     private ObjectHistogram histogram;
@@ -93,6 +95,9 @@
     @Before
     public void setUp() throws IOException {
         storage = mock(Storage.class);
+        add = mock(Add.class);
+        when(storage.createAdd(any(Category.class))).thenReturn(add);
+
         when(storage.getAgentId()).thenReturn("test");
         when(storage.createQuery()).then(new Answer<Query>() {
             @Override
@@ -177,6 +182,7 @@
         heapInfo = null;
         dao = null;
         storage = null;
+        add = null;
     }
 
     @Test
@@ -198,7 +204,9 @@
     public void testPutHeapInfo() throws IOException {
         dao.putHeapInfo(heapInfo, heapDumpData, histogram);
 
-        verify(storage).putPojo(HeapDAO.heapInfoCategory, false, heapInfo);
+        verify(storage).createAdd(HeapDAO.heapInfoCategory);
+        verify(add).setPojo(heapInfo);
+        verify(add).apply();
 
         ArgumentCaptor<InputStream> data = ArgumentCaptor.forClass(InputStream.class);
         verify(storage).saveFile(eq("heapdump-test-123-12345"), data.capture());
@@ -224,7 +232,9 @@
     public void testPutHeapInfoWithoutDump() throws IOException {
         dao.putHeapInfo(heapInfo, null, null);
 
-        verify(storage).putPojo(HeapDAO.heapInfoCategory, false, heapInfo);
+        verify(storage).createAdd(HeapDAO.heapInfoCategory);
+        verify(add).setPojo(heapInfo);
+        verify(add).apply();
 
         verify(storage, never()).saveFile(anyString(), any(InputStream.class));
         assertEquals("test-123-12345", heapInfo.getHeapId());
--- a/vm-memory/common/src/main/java/com/redhat/thermostat/vm/memory/common/internal/VmMemoryStatDAOImpl.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/vm-memory/common/src/main/java/com/redhat/thermostat/vm/memory/common/internal/VmMemoryStatDAOImpl.java	Thu Jan 10 14:53:32 2013 +0100
@@ -42,9 +42,10 @@
 import com.redhat.thermostat.common.dao.VmRef;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Query;
+import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Storage;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.model.VmMemoryStat;
 import com.redhat.thermostat.vm.memory.common.VmMemoryStatDAO;
 
@@ -76,7 +77,9 @@
 
     @Override
     public void putVmMemoryStat(VmMemoryStat stat) {
-        storage.putPojo(vmMemoryStatsCategory, false, stat);
+        Put add = storage.createAdd(vmMemoryStatsCategory);
+        add.setPojo(stat);
+        add.apply();
     }
 
     @Override
--- a/vm-memory/common/src/test/java/com/redhat/thermostat/vm/memory/common/internal/VmMemoryStatDAOTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/vm-memory/common/src/test/java/com/redhat/thermostat/vm/memory/common/internal/VmMemoryStatDAOTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -54,17 +54,18 @@
 
 import com.redhat.thermostat.common.dao.HostRef;
 import com.redhat.thermostat.common.dao.VmRef;
+import com.redhat.thermostat.storage.core.Add;
+import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
 import com.redhat.thermostat.storage.core.Query;
+import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Storage;
-import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.model.VmMemoryStat;
 import com.redhat.thermostat.storage.model.VmMemoryStat.Generation;
 import com.redhat.thermostat.storage.model.VmMemoryStat.Space;
 import com.redhat.thermostat.test.MockQuery;
 import com.redhat.thermostat.vm.memory.common.VmMemoryStatDAO;
-import com.redhat.thermostat.vm.memory.common.internal.VmMemoryStatDAOImpl;
 
 public class VmMemoryStatDAOTest {
 
@@ -194,9 +195,14 @@
         VmMemoryStat stat = new VmMemoryStat(1, 2, generations.toArray(new Generation[generations.size()]));
 
         Storage storage = mock(Storage.class);
+        Add add = mock(Add.class);
+        when(storage.createAdd(any(Category.class))).thenReturn(add);
+        
         VmMemoryStatDAO dao = new VmMemoryStatDAOImpl(storage);
         dao.putVmMemoryStat(stat);
 
-        verify(storage).putPojo(VmMemoryStatDAO.vmMemoryStatsCategory, false, stat);
+        verify(storage).createAdd(VmMemoryStatDAO.vmMemoryStatsCategory);
+        verify(add).setPojo(stat);
+        verify(add).apply();
     }
 }
--- a/web/client/src/main/java/com/redhat/thermostat/web/client/internal/WebStorage.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/web/client/src/main/java/com/redhat/thermostat/web/client/internal/WebStorage.java	Thu Jan 10 14:53:32 2013 +0100
@@ -43,6 +43,7 @@
 import java.io.OutputStream;
 import java.io.Reader;
 import java.lang.reflect.Array;
+import java.lang.reflect.InvocationTargetException;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.security.SecureRandom;
@@ -57,6 +58,7 @@
 
 import javax.net.ssl.SSLContext;
 
+import org.apache.commons.beanutils.BeanUtils;
 import org.apache.commons.codec.binary.Base64;
 import org.apache.http.Header;
 import org.apache.http.HttpEntity;
@@ -84,17 +86,20 @@
 import com.redhat.thermostat.common.ssl.SslInitException;
 import com.redhat.thermostat.common.utils.LoggingUtils;
 import com.redhat.thermostat.storage.config.StartupConfiguration;
+import com.redhat.thermostat.storage.core.Add;
 import com.redhat.thermostat.storage.core.AuthToken;
+import com.redhat.thermostat.storage.core.BasePut;
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Connection;
 import com.redhat.thermostat.storage.core.Cursor;
+import com.redhat.thermostat.storage.core.Key;
 import com.redhat.thermostat.storage.core.Query;
 import com.redhat.thermostat.storage.core.Remove;
+import com.redhat.thermostat.storage.core.Replace;
 import com.redhat.thermostat.storage.core.SecureStorage;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.core.StorageException;
 import com.redhat.thermostat.storage.core.Update;
-import com.redhat.thermostat.storage.model.AgentIdPojo;
 import com.redhat.thermostat.storage.model.Pojo;
 import com.redhat.thermostat.web.common.ThermostatGSONConverter;
 import com.redhat.thermostat.web.common.WebInsert;
@@ -271,6 +276,26 @@
 
     }
 
+    private class WebAdd extends BasePut implements Add {
+
+        @Override
+        public void apply() {
+            int categoryId = getCategoryId(getCategory());
+            putImpl(new WebInsert(categoryId, false, getPojo().getClass().getName()), getPojo());
+        }
+        
+    }
+
+    private class WebReplace extends BasePut implements Replace {
+
+        @Override
+        public void apply() {
+            int categoryId = getCategoryId(getCategory());
+            putImpl(new WebInsert(categoryId, true, getPojo().getClass().getName()), getPojo());
+        }
+        
+    }
+
     private String endpoint;
     private UUID agentId;
 
@@ -483,18 +508,22 @@
     }
 
     @Override
-    public void putPojo(final Category category, final boolean replace, final AgentIdPojo pojo)
-            throws StorageException {
+    public Add createAdd(Category into) {
+        WebAdd add = new WebAdd();
+        add.setCategory(into);
+        return add;
+    }
 
-        // TODO: This logic should probably be moved elsewhere. I.e. out of the
-        // Storage API.
-        if (pojo.getAgentId() == null) {
-            pojo.setAgentId(getAgentId());
-        }
+    @Override
+    public Replace createReplace(Category into) {
+        WebReplace replace = new WebReplace();
+        replace.setCategory(into);
+        return replace;
+    }
 
-        int categoryId = categoryIds.get(category);
-        WebInsert insert = new WebInsert(categoryId, replace, pojo.getClass()
-                .getName());
+    private void putImpl(WebInsert insert, final Pojo pojo) throws StorageException {
+
+        maybeAddAgentId(pojo);
         NameValuePair insertParam = new BasicNameValuePair("insert",
                 gson.toJson(insert));
         NameValuePair pojoParam = new BasicNameValuePair("pojo",
@@ -504,6 +533,16 @@
 
     }
 
+    private void maybeAddAgentId(final Pojo pojo) throws AssertionError {
+        try {
+            if (BeanUtils.getProperty(pojo, Key.AGENT_ID.getName()) == null) {
+                BeanUtils.setProperty(pojo, Key.AGENT_ID.getName(), getAgentId());
+            }
+        } catch (IllegalAccessException | InvocationTargetException | NoSuchMethodException e) {
+            throw new AssertionError("Pojo needs to have an agentId property");
+        }
+    }
+
     @Override
     public void removePojo(Remove remove) throws StorageException {
         NameValuePair removeParam = new BasicNameValuePair("remove",
@@ -599,4 +638,8 @@
         // Nothing to do here.
     }
 
+    int getCategoryId(Category category) {
+        return categoryIds.get(category);
+    }
+
 }
--- a/web/client/src/test/java/com/redhat/thermostat/web/client/internal/WebStorageTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/web/client/src/test/java/com/redhat/thermostat/web/client/internal/WebStorageTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -86,6 +86,7 @@
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Query;
 import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Remove;
@@ -289,7 +290,13 @@
         TestObj obj = new TestObj();
         obj.setProperty1("fluff");
 
-        storage.putPojo(category, true, obj);
+        // We need an agentId, so that we can check automatic insert of agentId.
+        UUID agentId = new UUID(1, 2);
+        storage.setAgentId(agentId);
+
+        Put replace = storage.createReplace(category);
+        replace.setPojo(obj);
+        replace.apply();
 
         Gson gson = new Gson();
         StringReader reader = new StringReader(requestBody);
@@ -308,6 +315,9 @@
         assertEquals(2, parts.length);
         assertEquals("pojo", parts[0]);
         Object resultObj = gson.fromJson(parts[1], Class.forName(insert.getPojoClass()));
+
+        // Set agentId on expected object, because we expect WebStorage to insert it for us.
+        obj.setAgentId(agentId.toString());
         assertEquals(obj, resultObj);
     }
 
--- a/web/common/src/main/java/com/redhat/thermostat/web/common/WebInsert.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/web/common/src/main/java/com/redhat/thermostat/web/common/WebInsert.java	Thu Jan 10 14:53:32 2013 +0100
@@ -45,7 +45,6 @@
     private String pojoClass;
 
     public WebInsert() {
-        this(-1, false, null);
     }
 
     public WebInsert(int categoryId, boolean replace, String pojoClass) {
--- a/web/server/src/main/java/com/redhat/thermostat/web/server/WebStorageEndPoint.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/web/server/src/main/java/com/redhat/thermostat/web/server/WebStorageEndPoint.java	Thu Jan 10 14:53:32 2013 +0100
@@ -65,12 +65,12 @@
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Query;
 import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Remove;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.core.Update;
-import com.redhat.thermostat.storage.model.AgentIdPojo;
 import com.redhat.thermostat.storage.model.Pojo;
 import com.redhat.thermostat.web.common.Qualifier;
 import com.redhat.thermostat.web.common.StorageWrapper;
@@ -250,12 +250,14 @@
         try {
             String insertParam = req.getParameter("insert");
             WebInsert insert = gson.fromJson(insertParam, WebInsert.class);
-            Class<? extends AgentIdPojo> pojoCls = (Class<? extends AgentIdPojo>) Class.forName(insert.getPojoClass());
+            Class<? extends Pojo> pojoCls = (Class<? extends Pojo>) Class.forName(insert.getPojoClass());
             String pojoParam = req.getParameter("pojo");
-            AgentIdPojo pojo = gson.fromJson(pojoParam, pojoCls);
+            Pojo pojo = gson.fromJson(pojoParam, pojoCls);
             int categoryId = insert.getCategoryId();
             Category category = getCategoryFromId(categoryId);
-            storage.putPojo(category, insert.isReplace(), pojo);
+            Put targetPut = insert.isReplace() ? storage.createReplace(category) : storage.createAdd(category);
+            targetPut.setPojo(pojo);
+            targetPut.apply();
             resp.setStatus(HttpServletResponse.SC_OK);
         } catch (ClassNotFoundException ex) {
             resp.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
--- a/web/server/src/test/java/com/redhat/thermostat/web/server/WebStorageEndpointTest.java	Fri Dec 21 14:19:06 2012 +0100
+++ b/web/server/src/test/java/com/redhat/thermostat/web/server/WebStorageEndpointTest.java	Thu Jan 10 14:53:32 2013 +0100
@@ -86,6 +86,7 @@
 import com.redhat.thermostat.storage.core.Query.Criteria;
 import com.redhat.thermostat.storage.core.Query.SortDirection;
 import com.redhat.thermostat.storage.core.Remove;
+import com.redhat.thermostat.storage.core.Replace;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.core.Update;
 import com.redhat.thermostat.storage.model.BasePojo;
@@ -286,6 +287,9 @@
     @Test
     public void testPutPojo() throws IOException {
 
+        Replace replace = mock(Replace.class);
+        when(mockStorage.createReplace(any(Category.class))).thenReturn(replace);
+
         TestClass expected1 = new TestClass();
         expected1.setKey1("fluff1");
         expected1.setKey2(42);
@@ -310,7 +314,9 @@
         out.write("\n");
         out.flush();
         assertEquals(200, conn.getResponseCode());
-        verify(mockStorage).putPojo(category, true, expected1);
+        verify(mockStorage).createReplace(category);
+        verify(replace).setPojo(expected1);
+        verify(replace).apply();
     }
 
     private void sendAuthorization(HttpURLConnection conn, String username, String passwd) {