changeset 1218:e2034aa58edf

Implement aggregate prepared query (count). Reviewed-by: omajid Review-thread: http://icedtea.classpath.org/pipermail/thermostat/2013-August/007834.html PR1509
author Severin Gehwolf <sgehwolf@redhat.com>
date Fri, 09 Aug 2013 10:43:36 +0200
parents b43db0f054d2
children 57db66363d91
files integration-tests/src/test/java/com/redhat/thermostat/itest/WebAppTest.java storage/core/src/main/java/com/redhat/thermostat/storage/core/AggregateQuery.java storage/core/src/main/java/com/redhat/thermostat/storage/core/BackingStorage.java storage/core/src/main/java/com/redhat/thermostat/storage/core/Category.java storage/core/src/main/java/com/redhat/thermostat/storage/core/CategoryAdapter.java storage/core/src/main/java/com/redhat/thermostat/storage/core/QueuedBackingStorage.java storage/core/src/main/java/com/redhat/thermostat/storage/dao/AgentInfoDAO.java storage/core/src/main/java/com/redhat/thermostat/storage/internal/AdaptedCategory.java storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/AgentInfoDAOImpl.java storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/BaseCountable.java storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/DAOImplStatementDescriptorRegistration.java storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/HostInfoDAOImpl.java storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/VmInfoDAOImpl.java storage/core/src/main/java/com/redhat/thermostat/storage/internal/statement/StatementDescriptorParser.java storage/core/src/main/java/com/redhat/thermostat/storage/model/AggregateCount.java storage/core/src/main/java/com/redhat/thermostat/storage/model/AggregateResult.java storage/core/src/test/java/com/redhat/thermostat/storage/core/CategoryAdapterTest.java storage/core/src/test/java/com/redhat/thermostat/storage/core/CategoryTest.java storage/core/src/test/java/com/redhat/thermostat/storage/internal/AdaptedCategoryTest.java storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/AgentInfoDAOTest.java storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/BaseCountableTest.java storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/DAOImplStatementDescriptorRegistrationTest.java storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/HostInfoDAOTest.java storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/VmInfoDAOTest.java storage/core/src/test/java/com/redhat/thermostat/storage/internal/statement/StatementDescriptorParserTest.java storage/core/src/test/java/com/redhat/thermostat/storage/model/AggregateCountTest.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 web/client/src/main/java/com/redhat/thermostat/web/client/internal/WebStorage.java web/common/src/test/java/com/redhat/thermostat/web/common/ThermostatGSONConverterTest.java web/server/src/main/java/com/redhat/thermostat/web/server/WebStorageEndPoint.java web/server/src/test/java/com/redhat/thermostat/web/server/KnownDescriptorRegistryTest.java web/server/src/test/java/com/redhat/thermostat/web/server/WebStorageEndpointTest.java
diffstat 33 files changed, 1407 insertions(+), 113 deletions(-) [+]
line wrap: on
line diff
--- a/integration-tests/src/test/java/com/redhat/thermostat/itest/WebAppTest.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/integration-tests/src/test/java/com/redhat/thermostat/itest/WebAppTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -75,6 +75,8 @@
 import com.redhat.thermostat.storage.config.ConnectionConfiguration;
 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.CategoryAdapter;
 import com.redhat.thermostat.storage.core.Connection.ConnectionListener;
 import com.redhat.thermostat.storage.core.Connection.ConnectionStatus;
 import com.redhat.thermostat.storage.core.Cursor;
@@ -85,6 +87,9 @@
 import com.redhat.thermostat.storage.core.StatementExecutionException;
 import com.redhat.thermostat.storage.core.Storage;
 import com.redhat.thermostat.storage.core.auth.DescriptorMetadata;
+import com.redhat.thermostat.storage.dao.HostInfoDAO;
+import com.redhat.thermostat.storage.model.AggregateCount;
+import com.redhat.thermostat.storage.model.HostInfo;
 import com.redhat.thermostat.test.FreePortFinder;
 import com.redhat.thermostat.test.FreePortFinder.TryPort;
 import com.redhat.thermostat.vm.classstat.common.VmClassStatDAO;
@@ -389,8 +394,37 @@
 
         storage.getConnection().disconnect();
     }
+    
+    private static void addHostInfoData(int numberOfItems) throws IOException {
+        String[] roleNames = new String[] {
+                Roles.REGISTER_CATEGORY,
+                Roles.ACCESS_REALM,
+                Roles.LOGIN,
+                Roles.APPEND
+        };
+        Storage storage = getAndConnectStorage(PREP_USER, PREP_PASSWORD, roleNames);
+        storage.registerCategory(HostInfoDAO.hostInfoCategory);
+
+        for (int i = 0; i < numberOfItems; i++) {
+            HostInfo hostInfo = new HostInfo("foo " + i, "linux " + i, "kernel", "t8", i, i * 1000);
+            hostInfo.setAgentId("test-host-agent-id");
+            Add add = storage.createAdd(HostInfoDAO.hostInfoCategory);
+            add.setPojo(hostInfo);
+            add.apply();
+        }
+
+        storage.getConnection().disconnect();
+    }
 
     private static void deleteCpuData() throws IOException {
+        doDeleteData(CpuStatDAO.cpuStatCategory, "test-agent-id");
+    }
+    
+    private static void deleteHostInfoData() throws IOException {
+        doDeleteData(HostInfoDAO.hostInfoCategory, "test-host-agent-id");
+    }
+    
+    private static void doDeleteData(Category<?> category, String agentId) throws IOException {
         String[] roleNames = new String[] {
                 Roles.REGISTER_CATEGORY,
                 Roles.ACCESS_REALM,
@@ -398,10 +432,8 @@
                 Roles.PURGE
         };
         Storage storage = getAndConnectStorage(PREP_USER, PREP_PASSWORD, roleNames);
-        storage.registerCategory(CpuStatDAO.cpuStatCategory);
-
-        storage.purge("test-agent-id");
-
+        storage.registerCategory(category);
+        storage.purge(agentId);
         storage.getConnection().disconnect();
     }
 
@@ -488,6 +520,44 @@
     }
     
     @Test
+    public void authorizedAggregateCount() throws Exception {
+        try {
+            int count = 2;
+            // registers host info category
+            addHostInfoData(count);
+            
+            String[] roleNames = new String[] {
+                    Roles.REGISTER_CATEGORY,
+                    Roles.READ,
+                    Roles.LOGIN,
+                    Roles.ACCESS_REALM,
+                    Roles.PREPARE_STATEMENT,
+                    Roles.GRANT_READ_ALL // don't want to test filtered results
+            };
+            Storage webStorage = getAndConnectStorage(TEST_USER, TEST_PASSWORD, roleNames);
+            Category<AggregateCount> adapted = new CategoryAdapter<HostInfo, AggregateCount>(HostInfoDAO.hostInfoCategory).getAdapted(AggregateCount.class);
+            // register adapted category.
+            webStorage.registerCategory(adapted);
+            
+            // storage-core registers this descriptor. no need to do it in this
+            // test.
+            String strDesc = "QUERY-COUNT host-info";
+            StatementDescriptor<AggregateCount> queryDesc = new StatementDescriptor<>(adapted, strDesc);
+            PreparedStatement<AggregateCount> query = webStorage.prepareStatement(queryDesc);
+    
+            Cursor<AggregateCount> cursor = query.executeQuery();
+            assertTrue(cursor.hasNext());
+            AggregateCount c = cursor.next();
+            assertFalse(cursor.hasNext());
+            assertEquals(count, c.getCount());
+    
+            webStorage.getConnection().disconnect();
+        } finally {
+            deleteHostInfoData();
+        }
+    }
+    
+    @Test
     public void authorizedQueryEqualTo() throws Exception {
 
         String[] roleNames = new String[] {
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/core/AggregateQuery.java	Fri Aug 09 10:43:36 2013 +0200
@@ -0,0 +1,91 @@
+/*
+ * Copyright 2012, 2013 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;
+import com.redhat.thermostat.storage.query.Expression;
+
+/**
+ * Common super class for aggregate queries.
+ */
+public abstract class AggregateQuery<T extends Pojo> implements Query<T> {
+    
+    public enum AggregateFunction {
+        /**
+         * Aggregate records by counting them.
+         */
+        COUNT
+    }
+    
+    protected final Query<T> queryToAggregate;
+    private final AggregateFunction function;
+
+    public AggregateQuery(AggregateFunction function, Query<T> queryToAggregate) {
+        this.function = function;
+        this.queryToAggregate = queryToAggregate;
+    }
+    
+    @Override
+    public void where(Expression expr) {
+        queryToAggregate.where(expr);
+    }
+
+    @Override
+    public void sort(Key<?> key,
+            SortDirection direction) {
+        queryToAggregate.sort(key, direction);
+    }
+
+    @Override
+    public void limit(int n) {
+        queryToAggregate.limit(n);
+    }
+
+    @Override
+    public Expression getWhereExpression() {
+        return queryToAggregate.getWhereExpression();
+    }
+    
+    /**
+     * 
+     * @return The function by which to aggregate by.
+     */
+    public AggregateFunction getAggregateFunction() {
+        return this.function;
+    }
+
+}
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/core/BackingStorage.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/core/BackingStorage.java	Fri Aug 09 10:43:36 2013 +0200
@@ -36,6 +36,7 @@
 
 package com.redhat.thermostat.storage.core;
 
+import com.redhat.thermostat.storage.core.AggregateQuery.AggregateFunction;
 import com.redhat.thermostat.storage.model.Pojo;
 
 /**
@@ -52,6 +53,8 @@
     
     <T extends Pojo> Query<T> createQuery(Category<T> category);
     
+    <T extends Pojo> Query<T> createAggregateQuery(AggregateFunction function, Category<T> category);
+    
     // TODO Move createUpdate and createRemove here
 
 }
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/core/Category.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/core/Category.java	Fri Aug 09 10:43:36 2013 +0200
@@ -49,15 +49,15 @@
  */
 public class Category<T extends Pojo> {
 
-    private String name;
-    private final Map<String, Key<?>> keys;
+    protected String name;
+    protected Map<String, Key<?>> keys;
     private transient Class<T> dataClass;
-    private String dataClassName;
+    protected String dataClassName;
 
     public Category() {
         this(null, null);
     }
-
+    
     /**
      * Creates a new Category instance with the specified name.
      *
@@ -124,10 +124,11 @@
 
     @Override
     public String toString() {
-        return getName() + keys;
+        return getName() + "|" + getDataClass().getName() + "|" + keys;
     }
+    
     public int hashCode() {
-        return Objects.hash(name, keys);
+        return Objects.hash(name, keys, getDataClass());
     }
 
     public boolean equals(Object o) {
@@ -135,7 +136,9 @@
             return false;
         }
         Category<?> other = (Category<?>) o;
-        return Objects.equals(name, other.name) && Objects.equals(keys, other.keys);
+        return Objects.equals(name, other.name) &&
+                Objects.equals(keys, other.keys) &&
+                Objects.equals(getDataClass(), other.getDataClass());
     }
 }
 
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/core/CategoryAdapter.java	Fri Aug 09 10:43:36 2013 +0200
@@ -0,0 +1,77 @@
+/*
+ * Copyright 2012, 2013 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 java.util.Objects;
+
+import com.redhat.thermostat.storage.internal.AdaptedCategory;
+import com.redhat.thermostat.storage.model.Pojo;
+
+/**
+ * 
+ * Adapts a given category to an aggregate equivalent.
+ *
+ * @param <T> The source data type.
+ * @param <S> The target data type after adaptation.
+ */
+public class CategoryAdapter<T extends Pojo, S extends Pojo> {
+
+    private final Category<T> sourceCategory;
+    
+    /**
+     * Constructor.
+     * 
+     * @param sourceCategory
+     *            A known source category.
+     * @throws NullPointerException
+     *             if sourceCategory was null.
+     * @throws IllegalArgumentException
+     *             if sourceCategory is not known.
+     */
+    public CategoryAdapter(Category<T> sourceCategory) {
+        Objects.requireNonNull(sourceCategory);
+        if (!Categories.contains(sourceCategory.getName())) {
+            throw new IllegalStateException("Only registered categories can be adapted!");
+        }
+        this.sourceCategory = sourceCategory;
+    }
+    
+    public Category<S> getAdapted(Class<S> targetType) {
+        AdaptedCategory<S, T> adapted = new AdaptedCategory<>(sourceCategory, targetType);
+        return adapted;
+    }
+}
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/core/QueuedBackingStorage.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/core/QueuedBackingStorage.java	Fri Aug 09 10:43:36 2013 +0200
@@ -38,6 +38,7 @@
 
 import java.util.concurrent.ExecutorService;
 
+import com.redhat.thermostat.storage.core.AggregateQuery.AggregateFunction;
 import com.redhat.thermostat.storage.model.Pojo;
 
 public class QueuedBackingStorage extends QueuedStorage implements
@@ -70,4 +71,10 @@
         return PreparedStatementFactory.getInstance(this, desc);
     }
 
+    @Override
+    public <T extends Pojo> Query<T> createAggregateQuery(
+            AggregateFunction function, Category<T> category) {
+        return ((BackingStorage) delegate).createAggregateQuery(function, category);
+    }
+
 }
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/dao/AgentInfoDAO.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/dao/AgentInfoDAO.java	Fri Aug 09 10:43:36 2013 +0200
@@ -62,7 +62,7 @@
             STOP_TIME_KEY,
             ALIVE_KEY,
             CONFIG_LISTEN_ADDRESS);
-
+    
     /**
      * Get information about all known agents.
      *
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/internal/AdaptedCategory.java	Fri Aug 09 10:43:36 2013 +0200
@@ -0,0 +1,76 @@
+/*
+ * Copyright 2012, 2013 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.internal;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import com.redhat.thermostat.storage.core.Category;
+import com.redhat.thermostat.storage.core.Key;
+import com.redhat.thermostat.storage.model.AggregateResult;
+import com.redhat.thermostat.storage.model.Pojo;
+
+/**
+ * An adapted category. This facilitates aggregate queries for which the data
+ * class type changes.
+ *
+ * @param <T> The type to adapt a category to.
+ * @param <S> The source type to adapt things from.
+ */
+public class AdaptedCategory<T extends Pojo, S extends Pojo> extends Category<T> {
+
+    /**
+     * Constructor used by CategoryAdapter which has just
+     * performed a registration check. That means only categories
+     * constructed via public Category constructors can get adapted.
+     *  
+     */
+    public AdaptedCategory(Category<S> category, Class<T> dataClass) {
+        this.name = category.getName();
+        Map<String, Key<?>> mappedKeys = new HashMap<>();
+        for (Key<?> key: category.getKeys()) {
+            mappedKeys.put(key.getName(), key);
+        }
+        this.keys = mappedKeys;
+        if (!AggregateResult.class.isAssignableFrom(dataClass)) {
+            String msg = "Can only adapt to aggregate results!";
+            throw new IllegalArgumentException(msg);
+        }
+        this.dataClassName = dataClass.getName();
+    }
+    
+}
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/AgentInfoDAOImpl.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/AgentInfoDAOImpl.java	Fri Aug 09 10:43:36 2013 +0200
@@ -43,6 +43,8 @@
 import java.util.logging.Logger;
 
 import com.redhat.thermostat.common.utils.LoggingUtils;
+import com.redhat.thermostat.storage.core.Category;
+import com.redhat.thermostat.storage.core.CategoryAdapter;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.DescriptorParsingException;
 import com.redhat.thermostat.storage.core.HostRef;
@@ -56,10 +58,11 @@
 import com.redhat.thermostat.storage.core.Update;
 import com.redhat.thermostat.storage.dao.AgentInfoDAO;
 import com.redhat.thermostat.storage.model.AgentInformation;
+import com.redhat.thermostat.storage.model.AggregateCount;
 import com.redhat.thermostat.storage.query.Expression;
 import com.redhat.thermostat.storage.query.ExpressionFactory;
 
-public class AgentInfoDAOImpl implements AgentInfoDAO {
+public class AgentInfoDAOImpl extends BaseCountable implements AgentInfoDAO {
 
     private static final Logger logger = LoggingUtils.getLogger(AgentInfoDAOImpl.class);
     static final String QUERY_AGENT_INFO = "QUERY "
@@ -67,37 +70,52 @@
             + Key.AGENT_ID.getName() + "' = ?s";
     static final String QUERY_ALL_AGENTS = "QUERY "
             + CATEGORY.getName();
+    // We can use AgentInfoDAO.CATEGORY.getName() here since this query
+    // only changes the data class. When executed we use the adapted
+    // aggregate category.
+    static final String AGGREGATE_COUNT_ALL_AGENTS = "QUERY-COUNT "
+            + CATEGORY.getName();
     static final String QUERY_ALIVE_AGENTS = "QUERY "
             + CATEGORY.getName() + " WHERE '" 
             + ALIVE_KEY.getName() + "' = ?b";
     
     private final Storage storage;
+    private final Category<AggregateCount> aggregateCategory;
     private final ExpressionFactory factory;
 
     public AgentInfoDAOImpl(Storage storage) {
         this.storage = storage;
+        // prepare adapted category and register it.
+        CategoryAdapter<AgentInformation, AggregateCount> adapter = new CategoryAdapter<>(CATEGORY);
+        this.aggregateCategory = adapter.getAdapted(AggregateCount.class);
         storage.registerCategory(CATEGORY);
+        storage.registerCategory(aggregateCategory);
         this.factory = new ExpressionFactory();
     }
 
     @Override
     public long getCount() {
-        long count = 0L;
-        Cursor<AgentInformation> agentCursor = getCursorForAllAgentInformation();
-        if (agentCursor == null) {
-            return count;
-        }
-        while (agentCursor.hasNext()) {
-            count++;
-            agentCursor.next();
-        }
+        StatementDescriptor<AggregateCount> desc = new StatementDescriptor<>(
+                aggregateCategory, AGGREGATE_COUNT_ALL_AGENTS);
+        long count = getCount(desc, storage);
         return count;
     }
 
     @Override
     public List<AgentInformation> getAllAgentInformation() {
-        Cursor<AgentInformation> agentCursor = getCursorForAllAgentInformation();
-        if (agentCursor == null) {
+        Cursor<AgentInformation> agentCursor;
+        StatementDescriptor<AgentInformation> desc = new StatementDescriptor<>(CATEGORY, QUERY_ALL_AGENTS);
+        PreparedStatement<AgentInformation> prepared = null;
+        try {
+            prepared = storage.prepareStatement(desc);
+            agentCursor =  prepared.executeQuery();
+        } catch (DescriptorParsingException e) {
+            // should not happen, but if it *does* happen, at least log it
+            logger.log(Level.SEVERE, "Preparing query '" + desc + "' failed!", e);
+            return Collections.emptyList();
+        } catch (StatementExecutionException e) {
+            // should not happen, but if it *does* happen, at least log it
+            logger.log(Level.SEVERE, "Executing query '" + desc + "' failed!", e);
             return Collections.emptyList();
         }
         List<AgentInformation> results = new ArrayList<>();
@@ -108,24 +126,6 @@
         }
         return results;
     }
-    
-    private Cursor<AgentInformation> getCursorForAllAgentInformation() {
-        StatementDescriptor<AgentInformation> desc = new StatementDescriptor<>(CATEGORY, QUERY_ALL_AGENTS);
-        PreparedStatement<AgentInformation> prepared = null;
-        try {
-            prepared = storage.prepareStatement(desc);
-            return prepared.executeQuery();
-        } catch (DescriptorParsingException e) {
-            // should not happen, but if it *does* happen, at least log it
-            logger.log(Level.SEVERE, "Preparing query '" + desc + "' failed!", e);
-            return null;
-        } catch (StatementExecutionException e) {
-            // should not happen, but if it *does* happen, at least log it
-            logger.log(Level.SEVERE, "Executing query '" + desc + "' failed!", e);
-            return null;
-        }
-        
-    }
 
     @Override
     public List<AgentInformation> getAliveAgents() {
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/BaseCountable.java	Fri Aug 09 10:43:36 2013 +0200
@@ -0,0 +1,86 @@
+/*
+ * Copyright 2012, 2013 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.internal.dao;
+
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+import com.redhat.thermostat.common.utils.LoggingUtils;
+import com.redhat.thermostat.storage.core.Cursor;
+import com.redhat.thermostat.storage.core.DescriptorParsingException;
+import com.redhat.thermostat.storage.core.PreparedStatement;
+import com.redhat.thermostat.storage.core.StatementDescriptor;
+import com.redhat.thermostat.storage.core.StatementExecutionException;
+import com.redhat.thermostat.storage.core.Storage;
+import com.redhat.thermostat.storage.model.AggregateCount;
+
+class BaseCountable {
+    
+    private static final int ERROR_COUNT_RESULT = -1;
+    private static final Logger logger = LoggingUtils.getLogger(BaseCountable.class);
+
+    /**
+     * Performs an aggregate count query as described by the given descriptor.
+     * 
+     * @param desc
+     *            The no-free-variables statement descriptor.
+     * @param storage
+     *            The storage to use for preparing the descriptor.
+     * 
+     * @return -1 if execution failed for some reason, the actual count of the
+     *         query results if successful.
+     */
+    protected long getCount(StatementDescriptor<AggregateCount> desc, Storage storage) {
+        PreparedStatement<AggregateCount> prepared = null;
+        Cursor<AggregateCount> countCursor = null;
+        try {
+            prepared = storage.prepareStatement(desc);
+            countCursor =  prepared.executeQuery();
+        } catch (DescriptorParsingException e) {
+            // should not happen, but if it *does* happen, at least log it
+            logger.log(Level.SEVERE, "Preparing query '" + desc + "' failed!", e);
+            return ERROR_COUNT_RESULT;
+        } catch (StatementExecutionException e) {
+            // should not happen, but if it *does* happen, at least log it
+            logger.log(Level.SEVERE, "Executing query '" + desc + "' failed!", e);
+            return ERROR_COUNT_RESULT;
+        }
+        // there is only one result
+        AggregateCount aggreateResult = countCursor.next();
+        return aggreateResult.getCount();
+    }
+}
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/DAOImplStatementDescriptorRegistration.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/DAOImplStatementDescriptorRegistration.java	Fri Aug 09 10:43:36 2013 +0200
@@ -58,13 +58,16 @@
         daoDescs.add(AgentInfoDAOImpl.QUERY_AGENT_INFO);
         daoDescs.add(AgentInfoDAOImpl.QUERY_ALIVE_AGENTS);
         daoDescs.add(AgentInfoDAOImpl.QUERY_ALL_AGENTS);
+        daoDescs.add(AgentInfoDAOImpl.AGGREGATE_COUNT_ALL_AGENTS);
         daoDescs.add(BackendInfoDAOImpl.QUERY_BACKEND_INFO);
         daoDescs.add(HostInfoDAOImpl.QUERY_HOST_INFO);
         daoDescs.add(HostInfoDAOImpl.QUERY_ALL_HOSTS);
+        daoDescs.add(HostInfoDAOImpl.AGGREGATE_COUNT_ALL_HOSTS);
         daoDescs.add(NetworkInterfaceInfoDAOImpl.QUERY_NETWORK_INFO);
         daoDescs.add(VmInfoDAOImpl.QUERY_ALL_VMS_FOR_HOST);
         daoDescs.add(VmInfoDAOImpl.QUERY_ALL_VMS);
         daoDescs.add(VmInfoDAOImpl.QUERY_VM_INFO);
+        daoDescs.add(VmInfoDAOImpl.AGGREGATE_COUNT_ALL_VMS);
         return daoDescs;
     }
 
@@ -92,6 +95,9 @@
         } else if (descriptor.equals(HostInfoDAOImpl.QUERY_ALL_HOSTS)) {
             DescriptorMetadata metadata = new DescriptorMetadata();
             return metadata;
+        } else if (descriptor.equals(HostInfoDAOImpl.AGGREGATE_COUNT_ALL_HOSTS)) {
+            DescriptorMetadata metadata = new DescriptorMetadata();
+            return metadata;
         } else if (descriptor.equals(NetworkInterfaceInfoDAOImpl.QUERY_NETWORK_INFO)) {
             String agentId = (String)params[0].getValue();
             DescriptorMetadata metadata = new DescriptorMetadata(agentId);
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/HostInfoDAOImpl.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/HostInfoDAOImpl.java	Fri Aug 09 10:43:36 2013 +0200
@@ -44,6 +44,8 @@
 import java.util.logging.Logger;
 
 import com.redhat.thermostat.common.utils.LoggingUtils;
+import com.redhat.thermostat.storage.core.Category;
+import com.redhat.thermostat.storage.core.CategoryAdapter;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.DescriptorParsingException;
 import com.redhat.thermostat.storage.core.HostRef;
@@ -56,24 +58,34 @@
 import com.redhat.thermostat.storage.dao.AgentInfoDAO;
 import com.redhat.thermostat.storage.dao.HostInfoDAO;
 import com.redhat.thermostat.storage.model.AgentInformation;
+import com.redhat.thermostat.storage.model.AggregateCount;
 import com.redhat.thermostat.storage.model.HostInfo;
 
-public class HostInfoDAOImpl implements HostInfoDAO {
+public class HostInfoDAOImpl extends BaseCountable implements HostInfoDAO {
     
     private static final Logger logger = LoggingUtils.getLogger(HostInfoDAOImpl.class);
     static final String QUERY_HOST_INFO = "QUERY "
             + hostInfoCategory.getName() + " WHERE '" 
             + Key.AGENT_ID.getName() + "' = ?s LIMIT 1";
     static final String QUERY_ALL_HOSTS = "QUERY " + hostInfoCategory.getName();
+    // We can use hostInfoCategory.getName() here since this query
+    // only changes the data class. When executed we use the adapted
+    // aggregate category.
+    static final String AGGREGATE_COUNT_ALL_HOSTS = "QUERY-COUNT " + hostInfoCategory.getName();
 
     private final Storage storage;
     private final AgentInfoDAO agentInfoDao;
-
+    private final Category<AggregateCount> aggregateCategory;
+    
 
     public HostInfoDAOImpl(Storage storage, AgentInfoDAO agentInfo) {
         this.storage = storage;
         this.agentInfoDao = agentInfo;
+        // Adapt category to the aggregate form
+        CategoryAdapter<HostInfo, AggregateCount> adapter = new CategoryAdapter<>(hostInfoCategory);
+        this.aggregateCategory = adapter.getAdapted(AggregateCount.class);
         storage.registerCategory(hostInfoCategory);
+        storage.registerCategory(aggregateCategory);
     }
 
     @Override
@@ -168,15 +180,9 @@
 
     @Override
     public long getCount() {
-        long count = 0;
-        Cursor<HostInfo> hostInfoCursor = getAllHostInfoCursor();
-        if (hostInfoCursor == null) {
-            return count;
-        }
-        while (hostInfoCursor.hasNext()) {
-            count++;
-            hostInfoCursor.next();
-        }
+        StatementDescriptor<AggregateCount> desc = new StatementDescriptor<>(
+                aggregateCategory, AGGREGATE_COUNT_ALL_HOSTS);
+        long count = getCount(desc, storage);
         return count;
     }
     
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/VmInfoDAOImpl.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/internal/dao/VmInfoDAOImpl.java	Fri Aug 09 10:43:36 2013 +0200
@@ -44,6 +44,8 @@
 import java.util.logging.Logger;
 
 import com.redhat.thermostat.common.utils.LoggingUtils;
+import com.redhat.thermostat.storage.core.Category;
+import com.redhat.thermostat.storage.core.CategoryAdapter;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.DescriptorParsingException;
 import com.redhat.thermostat.storage.core.HostRef;
@@ -57,11 +59,12 @@
 import com.redhat.thermostat.storage.core.VmRef;
 import com.redhat.thermostat.storage.dao.DAOException;
 import com.redhat.thermostat.storage.dao.VmInfoDAO;
+import com.redhat.thermostat.storage.model.AggregateCount;
 import com.redhat.thermostat.storage.model.VmInfo;
 import com.redhat.thermostat.storage.query.Expression;
 import com.redhat.thermostat.storage.query.ExpressionFactory;
 
-public class VmInfoDAOImpl implements VmInfoDAO {
+public class VmInfoDAOImpl extends BaseCountable implements VmInfoDAO {
     
     private final Logger logger = LoggingUtils.getLogger(VmInfoDAOImpl.class);
     static final String QUERY_VM_INFO = "QUERY " 
@@ -72,13 +75,19 @@
             + vmInfoCategory.getName() + " WHERE '" 
             + Key.AGENT_ID.getName() + "' = ?s";
     static final String QUERY_ALL_VMS = "QUERY " + vmInfoCategory.getName();
+    static final String AGGREGATE_COUNT_ALL_VMS = "QUERY-COUNT " + vmInfoCategory.getName();
     
     private final Storage storage;
     private final ExpressionFactory factory;
+    private final Category<AggregateCount> aggregateCategory;
 
     public VmInfoDAOImpl(Storage storage) {
         this.storage = storage;
+        // Adapt category to the aggregate form
+        CategoryAdapter<VmInfo, AggregateCount> adapter = new CategoryAdapter<>(vmInfoCategory);
+        this.aggregateCategory = adapter.getAdapted(AggregateCount.class);
         storage.registerCategory(vmInfoCategory);
+        storage.registerCategory(aggregateCategory);
         factory = new ExpressionFactory();
     }
 
@@ -156,26 +165,9 @@
 
     @Override
     public long getCount() {
-        long count = 0;
-        StatementDescriptor<VmInfo> desc = new StatementDescriptor<>(vmInfoCategory, QUERY_ALL_VMS);
-        PreparedStatement<VmInfo> stmt;
-        Cursor<VmInfo> cursor;
-        try {
-            stmt = storage.prepareStatement(desc);
-            cursor = stmt.executeQuery();
-        } catch (DescriptorParsingException e) {
-            // should not happen, but if it *does* happen, at least log it
-            logger.log(Level.SEVERE, "Preparing query '" + desc + "' failed!", e);
-            return count;
-        } catch (StatementExecutionException e) {
-            // should not happen, but if it *does* happen, at least log it
-            logger.log(Level.SEVERE, "Executing query '" + desc + "' failed!", e);
-            return count;
-        }
-        while (cursor.hasNext()) {
-            count++;
-            cursor.next();
-        }
+        StatementDescriptor<AggregateCount> desc = new StatementDescriptor<>(
+                aggregateCategory, AGGREGATE_COUNT_ALL_VMS);
+        long count = getCount(desc, storage);
         return count;
     }
 
--- a/storage/core/src/main/java/com/redhat/thermostat/storage/internal/statement/StatementDescriptorParser.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/internal/statement/StatementDescriptorParser.java	Fri Aug 09 10:43:36 2013 +0200
@@ -40,6 +40,7 @@
 import java.util.List;
 import java.util.StringTokenizer;
 
+import com.redhat.thermostat.storage.core.AggregateQuery.AggregateFunction;
 import com.redhat.thermostat.storage.core.BackingStorage;
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.DescriptorParsingException;
@@ -61,7 +62,7 @@
  * 
  * <pre>
  * statementDesc := statementType category suffix
- * statementType := 'QUERY'
+ * statementType := 'QUERY' | 'QUERY-COUNT'
  * category      := literal
  * suffix        := 'WHERE' where |
  *                  'SORT' sortCond |
@@ -105,7 +106,7 @@
 
     private static final String TOKEN_DELIMS = " \t\r\n\f";
     private static final String[] KNOWN_STATEMENT_TYPES = new String[] {
-        "QUERY",
+        "QUERY", "QUERY-COUNT"
     };
     private static final String SORTLIST_SEP = ",";
     private static final String KEYWORD_WHERE = "WHERE";
@@ -579,10 +580,15 @@
 
     private void createStatement() {
         if (tokens[0].equals(KNOWN_STATEMENT_TYPES[0])) {
-            // query case
+            // regular query case
             Query<T> query = storage.createQuery(desc.getCategory());
             this.parsedStatement = new ParsedStatementImpl<>(query);
-        } else {
+        } else if (tokens[0].equals(KNOWN_STATEMENT_TYPES[1])) {
+            // create aggregate count query
+            Query<T> query = storage.createAggregateQuery(AggregateFunction.COUNT, desc.getCategory());
+            this.parsedStatement = new ParsedStatementImpl<>(query);
+        }
+        else {
             throw new IllegalStateException("Don't know how to create statement type '" + tokens[0] + "'");
         }
     }
@@ -603,9 +609,11 @@
     }
 
     private void matchStatementType() throws DescriptorParsingException {
-        // matches 'QUERY' only at this point
+        // matches 'QUERY' and 'QUERY-COUNT' only at this point
         if (tokens[currTokenIndex].equals(KNOWN_STATEMENT_TYPES[0])) {
             currTokenIndex++;
+        } else if (tokens[currTokenIndex].equals(KNOWN_STATEMENT_TYPES[1])) {
+            currTokenIndex++;
         } else {
             throw new DescriptorParsingException("Unknown statement type: '" + tokens[currTokenIndex] + "'");
         }
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/model/AggregateCount.java	Fri Aug 09 10:43:36 2013 +0200
@@ -0,0 +1,109 @@
+/*
+ * Copyright 2012, 2013 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;
+
+import java.util.NoSuchElementException;
+import java.util.Objects;
+
+import com.redhat.thermostat.storage.core.Cursor;
+import com.redhat.thermostat.storage.core.Entity;
+import com.redhat.thermostat.storage.core.Persist;
+
+/**
+ * Model class for aggregate counts.
+ *
+ */
+@Entity
+public class AggregateCount implements AggregateResult {
+
+    private long count;
+    
+    @Persist
+    public long getCount() {
+        return count;
+    }
+    
+    @Persist
+    public void setCount(long count) {
+        this.count = count;
+    }
+    
+    @Override
+    public boolean equals(Object other) {
+        if (!(other instanceof AggregateCount)) {
+            return false;
+        }
+        AggregateCount o = (AggregateCount)other;
+        return this.getCount() == o.getCount();
+    }
+    
+    @Override
+    public int hashCode() {
+        return Objects.hash(getCount());
+    }
+    
+    @SuppressWarnings("unchecked")
+    public <T extends Pojo> Cursor<T> getCursor() {
+        return (Cursor<T>) new AggregateCursor<>(this);
+    }
+    
+    private static class AggregateCursor<T extends Pojo> implements Cursor<T> {
+
+        private boolean available = true;
+        private final T count;
+        
+        private AggregateCursor(T count) {
+            this.count = count;
+        }
+        
+        @Override
+        public boolean hasNext() {
+            return available; 
+        }
+
+        @Override
+        public T next() {
+            if (available) {
+                available = false;
+                return count;
+            } else {
+                throw new NoSuchElementException();
+            }
+        }
+        
+    }
+}
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/main/java/com/redhat/thermostat/storage/model/AggregateResult.java	Fri Aug 09 10:43:36 2013 +0200
@@ -0,0 +1,46 @@
+/*
+ * Copyright 2012, 2013 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;
+
+/**
+ * Super type for aggregate results.
+ * Marker only.
+ *
+ */
+public interface AggregateResult extends Pojo {
+
+}
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/test/java/com/redhat/thermostat/storage/core/CategoryAdapterTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -0,0 +1,80 @@
+/*
+ * Copyright 2012, 2013 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.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import org.junit.Test;
+
+import com.redhat.thermostat.storage.dao.AgentInfoDAO;
+import com.redhat.thermostat.storage.model.AgentInformation;
+import com.redhat.thermostat.storage.model.AggregateCount;
+
+public class CategoryAdapterTest {
+
+    @Test
+    public void canAdaptToAggregateResultDataClass() {
+        CategoryAdapter<AgentInformation, AggregateCount> adapter = new CategoryAdapter<>(AgentInfoDAO.CATEGORY);
+        Category<AggregateCount> aggregateCountCat = adapter.getAdapted(AggregateCount.class);
+        assertEquals(AggregateCount.class, aggregateCountCat.getDataClass());
+        assertEquals(AgentInfoDAO.CATEGORY.getName(), aggregateCountCat.getName());
+        assertFalse(AgentInfoDAO.CATEGORY.equals(aggregateCountCat));
+    }
+    
+    @Test
+    public void canCreateAdapterFromNull() {
+        try {
+            new CategoryAdapter<>(null);
+        } catch (NullPointerException e) {
+            // pass
+        }
+    }
+    
+    @Test
+    public void cannotAdaptUnknown() {
+        Category<?> unknown = mock(Category.class);
+        when(unknown.getName()).thenReturn("foo");
+        try {
+            new CategoryAdapter<>(unknown);
+        } catch (IllegalStateException e) {
+            // pass
+        }
+    }
+}
--- a/storage/core/src/test/java/com/redhat/thermostat/storage/core/CategoryTest.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/test/java/com/redhat/thermostat/storage/core/CategoryTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -37,13 +37,18 @@
 package com.redhat.thermostat.storage.core;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
 import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
 
 import org.junit.Test;
 
+import com.redhat.thermostat.storage.dao.HostInfoDAO;
 import com.redhat.thermostat.storage.model.Pojo;
 
 public class CategoryTest {
@@ -90,7 +95,26 @@
         Collection<Key<?>> keys = category.getKeys();
 
         keys.remove(key1);
-
+    }
+    
+    @Test
+    public void testEquals() {
+        Key<String> key1 = new Key<String>("key1", false);
+        Key<String> key2 = new Key<String>("key2", false);
+        Key<String> key3 = new Key<String>("key3", false);
+        Category<TestObj> category = new Category<>("testEquals", TestObj.class, key1, key2, key3);
+        assertTrue(category.equals(category));
+        assertFalse(category.equals(HostInfoDAO.hostInfoCategory));
+    }
+    
+    @Test
+    public void testHashCode() {
+        Key<String> key1 = new Key<String>("key1", false);
+        Category<TestObj> category = new Category<>("testHashCode", TestObj.class, key1);
+        Map<String, Key<?>> keys = new HashMap<>();
+        keys.put(key1.getName(), key1);
+        int expectedHash = Objects.hash("testHashCode", keys, TestObj.class);
+        assertEquals(expectedHash, category.hashCode());
     }
 }
 
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/test/java/com/redhat/thermostat/storage/internal/AdaptedCategoryTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -0,0 +1,82 @@
+/*
+ * Copyright 2012, 2013 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.internal;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import org.junit.Test;
+
+import com.redhat.thermostat.storage.dao.AgentInfoDAO;
+import com.redhat.thermostat.storage.dao.VmInfoDAO;
+import com.redhat.thermostat.storage.internal.AdaptedCategory;
+import com.redhat.thermostat.storage.model.AgentInformation;
+import com.redhat.thermostat.storage.model.AggregateCount;
+import com.redhat.thermostat.storage.model.AggregateResult;
+import com.redhat.thermostat.storage.model.VmInfo;
+
+public class AdaptedCategoryTest {
+
+    @Test
+    public void testEquals() {
+        AdaptedCategory<AggregateCount, AgentInformation> cat = new AdaptedCategory<>(AgentInfoDAO.CATEGORY, AggregateCount.class);
+        assertFalse(cat.equals(AgentInfoDAO.CATEGORY));
+        // equals self
+        assertEquals(cat, cat);
+        // not equal to any other category
+        assertFalse(VmInfoDAO.vmInfoCategory.equals(cat));
+    }
+    
+    @Test
+    public void getDataClass() {
+        AdaptedCategory<AggregateCount, AgentInformation> cat = new AdaptedCategory<>(AgentInfoDAO.CATEGORY, AggregateCount.class);
+        assertEquals(AggregateCount.class, cat.getDataClass());
+        assertTrue(AggregateResult.class.isAssignableFrom(cat.getDataClass()));
+    }
+    
+    @Test
+    public void adaptNonAggregateDataClass() {
+        try {
+            new AdaptedCategory<>(AgentInfoDAO.CATEGORY, VmInfo.class);
+        } catch (IllegalArgumentException e) {
+            // pass
+            assertTrue(e.getMessage().contains("Can only adapt to aggregate results"));
+        }
+    }
+    
+}
--- a/storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/AgentInfoDAOTest.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/AgentInfoDAOTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -66,6 +66,7 @@
 import com.redhat.thermostat.storage.core.Update;
 import com.redhat.thermostat.storage.dao.AgentInfoDAO;
 import com.redhat.thermostat.storage.model.AgentInformation;
+import com.redhat.thermostat.storage.model.AggregateCount;
 import com.redhat.thermostat.storage.query.Expression;
 import com.redhat.thermostat.storage.query.ExpressionFactory;
 
@@ -153,29 +154,26 @@
         assertEquals(expected, result);
     }
     
-    /*
-     * getCount() with two AgentInformation records.
-     */
     @Test
     public void testGetCount()
             throws DescriptorParsingException, StatementExecutionException {
-        AgentInformation agent2 = new AgentInformation();
+        AggregateCount count = new AggregateCount();
+        count.setCount(2);
         
         @SuppressWarnings("unchecked")
-        Cursor<AgentInformation> agentCursor = (Cursor<AgentInformation>) mock(Cursor.class);
-        when(agentCursor.hasNext()).thenReturn(true).thenReturn(true).thenReturn(false);
-        when(agentCursor.next()).thenReturn(agent1).thenReturn(agent2).thenReturn(null);
+        Cursor<AggregateCount> countCursor = (Cursor<AggregateCount>) mock(Cursor.class);
+        when(countCursor.next()).thenReturn(count);
 
         Storage storage = mock(Storage.class);
         @SuppressWarnings("unchecked")
-        PreparedStatement<AgentInformation> stmt = (PreparedStatement<AgentInformation>) mock(PreparedStatement.class);
-        when(storage.prepareStatement(anyDescriptor())).thenReturn(stmt);
-        when(stmt.executeQuery()).thenReturn(agentCursor);
+        PreparedStatement<AggregateCount> stmt = (PreparedStatement<AggregateCount>) mock(PreparedStatement.class);
+        @SuppressWarnings("unchecked")
+        StatementDescriptor<AggregateCount> desc = any(StatementDescriptor.class);
+        when(storage.prepareStatement(desc)).thenReturn(stmt);
+        when(stmt.executeQuery()).thenReturn(countCursor);
         AgentInfoDAOImpl dao = new AgentInfoDAOImpl(storage);
 
-        long count = dao.getCount();
-
-        assertEquals(2, count);
+        assertEquals(2, dao.getCount());
     }
     
     @SuppressWarnings("unchecked")
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/BaseCountableTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -0,0 +1,105 @@
+/*
+ * Copyright 2012, 2013 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.internal.dao;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.doThrow;
+
+import org.junit.Test;
+
+import com.redhat.thermostat.storage.core.Category;
+import com.redhat.thermostat.storage.core.Cursor;
+import com.redhat.thermostat.storage.core.DescriptorParsingException;
+import com.redhat.thermostat.storage.core.PreparedStatement;
+import com.redhat.thermostat.storage.core.StatementDescriptor;
+import com.redhat.thermostat.storage.core.StatementExecutionException;
+import com.redhat.thermostat.storage.core.Storage;
+import com.redhat.thermostat.storage.model.AggregateCount;
+
+public class BaseCountableTest {
+
+    @Test
+    public void testGetCountSuccessful() throws DescriptorParsingException, StatementExecutionException {
+        Storage storage = mock(Storage.class);
+        @SuppressWarnings("unchecked")
+        Category<AggregateCount> mockCategory = (Category<AggregateCount>)mock(Category.class);
+        BaseCountable countable = new BaseCountable();
+        String strDesc = "QUERY-COUNT vm-info";
+        StatementDescriptor<AggregateCount> desc = new StatementDescriptor<>(mockCategory, strDesc);
+        @SuppressWarnings("unchecked")
+        PreparedStatement<AggregateCount> prepared = (PreparedStatement<AggregateCount>)mock(PreparedStatement.class);
+        when(storage.prepareStatement(eq(desc))).thenReturn(prepared);
+        AggregateCount c = new AggregateCount();
+        c.setCount(3);
+        Cursor<AggregateCount> cursor = c.getCursor();
+        when(prepared.executeQuery()).thenReturn(cursor);
+        long count = countable.getCount(desc, storage);
+        assertEquals(3, count);
+    }
+    
+    @Test
+    public void testGetCountError() throws DescriptorParsingException, StatementExecutionException {
+        Storage storage = mock(Storage.class);
+        @SuppressWarnings("unchecked")
+        Category<AggregateCount> mockCategory = (Category<AggregateCount>)mock(Category.class);
+        BaseCountable countable = new BaseCountable();
+        String strDesc = "QUERY-COUNT vm-info";
+        StatementDescriptor<AggregateCount> desc = new StatementDescriptor<>(mockCategory, strDesc);
+        doThrow(DescriptorParsingException.class).when(storage).prepareStatement(eq(desc));
+        long count = countable.getCount(desc, storage);
+        assertEquals(-1, count);
+    }
+    
+    @Test
+    public void testGetCountError2() throws DescriptorParsingException, StatementExecutionException {
+        Storage storage = mock(Storage.class);
+        @SuppressWarnings("unchecked")
+        Category<AggregateCount> mockCategory = (Category<AggregateCount>)mock(Category.class);
+        BaseCountable countable = new BaseCountable();
+        String strDesc = "QUERY-COUNT vm-info";
+        StatementDescriptor<AggregateCount> desc = new StatementDescriptor<>(mockCategory, strDesc);
+        @SuppressWarnings("unchecked")
+        PreparedStatement<AggregateCount> prepared = (PreparedStatement<AggregateCount>)mock(PreparedStatement.class);
+        when(storage.prepareStatement(eq(desc))).thenReturn(prepared);
+        doThrow(StatementExecutionException.class).when(prepared).executeQuery();
+        long count = countable.getCount(desc, storage);
+        assertEquals(-1, count);
+    }
+}
--- a/storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/DAOImplStatementDescriptorRegistrationTest.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/DAOImplStatementDescriptorRegistrationTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -62,7 +62,7 @@
     public void registersAllQueries() {
         DAOImplStatementDescriptorRegistration reg = new DAOImplStatementDescriptorRegistration();
         Set<String> descriptors = reg.getStatementDescriptors();
-        assertEquals(10, descriptors.size());
+        assertEquals(13, descriptors.size());
         assertFalse(descriptors.contains(null));
     }
     
@@ -80,7 +80,7 @@
             registrations.add(r);
         }
         assertEquals(1, registrations.size());
-        assertEquals(10, registrations.get(0).getStatementDescriptors().size());
+        assertEquals(13, registrations.get(0).getStatementDescriptors().size());
     }
     
     @Test
--- a/storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/HostInfoDAOTest.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/HostInfoDAOTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -63,6 +63,7 @@
 import com.redhat.thermostat.storage.dao.AgentInfoDAO;
 import com.redhat.thermostat.storage.dao.HostInfoDAO;
 import com.redhat.thermostat.storage.model.AgentInformation;
+import com.redhat.thermostat.storage.model.AggregateCount;
 import com.redhat.thermostat.storage.model.HostInfo;
 
 
@@ -224,13 +225,25 @@
     }
 
     @Test
-    public void testGetCount() throws Exception {
-        Storage storage = setupStorageForSingleHost();
-        AgentInfoDAO agentInfo = mock(AgentInfoDAO.class);
+    public void testGetCount() throws DescriptorParsingException,
+            StatementExecutionException {
+        AggregateCount count = new AggregateCount();
+        count.setCount(2);
+
+        @SuppressWarnings("unchecked")
+        Cursor<AggregateCount> countCursor = (Cursor<AggregateCount>) mock(Cursor.class);
+        when(countCursor.next()).thenReturn(count);
 
-        HostInfoDAO hostsDAO = new HostInfoDAOImpl(storage, agentInfo);
+        Storage storage = mock(Storage.class);
+        @SuppressWarnings("unchecked")
+        PreparedStatement<AggregateCount> stmt = (PreparedStatement<AggregateCount>) mock(PreparedStatement.class);
+        @SuppressWarnings("unchecked")
+        StatementDescriptor<AggregateCount> desc = any(StatementDescriptor.class);
+        when(storage.prepareStatement(desc)).thenReturn(stmt);
+        when(stmt.executeQuery()).thenReturn(countCursor);
+        HostInfoDAOImpl dao = new HostInfoDAOImpl(storage, null);
 
-        assertEquals(1, hostsDAO.getCount());
+        assertEquals(2, dao.getCount());
     }
     
     @Test
--- a/storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/VmInfoDAOTest.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/test/java/com/redhat/thermostat/storage/internal/dao/VmInfoDAOTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -67,6 +67,7 @@
 import com.redhat.thermostat.storage.core.VmRef;
 import com.redhat.thermostat.storage.dao.DAOException;
 import com.redhat.thermostat.storage.dao.VmInfoDAO;
+import com.redhat.thermostat.storage.model.AggregateCount;
 import com.redhat.thermostat.storage.model.VmInfo;
 import com.redhat.thermostat.storage.query.ExpressionFactory;
 
@@ -287,9 +288,24 @@
     }
 
     @Test
-    public void testGetCount() throws Exception {
-        Storage storage = setupStorageForMultiVM();
-        VmInfoDAO dao = new VmInfoDAOImpl(storage);
+    public void testGetCount()
+            throws DescriptorParsingException, StatementExecutionException {
+        AggregateCount count = new AggregateCount();
+        count.setCount(2);
+        
+        @SuppressWarnings("unchecked")
+        Cursor<AggregateCount> countCursor = (Cursor<AggregateCount>) mock(Cursor.class);
+        when(countCursor.next()).thenReturn(count);
+
+        Storage storage = mock(Storage.class);
+        @SuppressWarnings("unchecked")
+        PreparedStatement<AggregateCount> stmt = (PreparedStatement<AggregateCount>) mock(PreparedStatement.class);
+        @SuppressWarnings("unchecked")
+        StatementDescriptor<AggregateCount> desc = any(StatementDescriptor.class);
+        when(storage.prepareStatement(desc)).thenReturn(stmt);
+        when(stmt.executeQuery()).thenReturn(countCursor);
+        VmInfoDAOImpl dao = new VmInfoDAOImpl(storage);
+
         assertEquals(2, dao.getCount());
     }
 
--- a/storage/core/src/test/java/com/redhat/thermostat/storage/internal/statement/StatementDescriptorParserTest.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/core/src/test/java/com/redhat/thermostat/storage/internal/statement/StatementDescriptorParserTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -42,6 +42,7 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -50,15 +51,21 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 
+import com.redhat.thermostat.storage.core.AggregateQuery;
 import com.redhat.thermostat.storage.core.BackingStorage;
+import com.redhat.thermostat.storage.core.Category;
+import com.redhat.thermostat.storage.core.CategoryAdapter;
 import com.redhat.thermostat.storage.core.DescriptorParsingException;
 import com.redhat.thermostat.storage.core.Key;
 import com.redhat.thermostat.storage.core.Query;
+import com.redhat.thermostat.storage.core.AggregateQuery.AggregateFunction;
 import com.redhat.thermostat.storage.core.Query.SortDirection;
 import com.redhat.thermostat.storage.core.StatementDescriptor;
 import com.redhat.thermostat.storage.dao.AgentInfoDAO;
 import com.redhat.thermostat.storage.model.AgentInformation;
+import com.redhat.thermostat.storage.model.AggregateCount;
 import com.redhat.thermostat.storage.internal.statement.BinaryExpressionNode;
 import com.redhat.thermostat.storage.internal.statement.LimitExpression;
 import com.redhat.thermostat.storage.internal.statement.NotBooleanExpressionNode;
@@ -95,6 +102,32 @@
         mockQuery = null;
     }
     
+    @SuppressWarnings({ "rawtypes", "unchecked" })
+    @Test
+    public void testParseAggregateCount() throws DescriptorParsingException {
+        Query<AggregateCount> query = mock(AggregateQuery.class);
+        ArgumentCaptor<Category> captor = ArgumentCaptor.forClass(Category.class);
+        when(storage.createAggregateQuery(eq(AggregateFunction.COUNT), captor.capture())).thenReturn(query);
+        // first adapt from the target category in order to be able to produce the
+        // right aggregate query with a different result type.
+        CategoryAdapter<AgentInformation, AggregateCount> adapter = new CategoryAdapter<>(AgentInfoDAO.CATEGORY);
+        Category<AggregateCount> aggregateCategory = adapter.getAdapted(AggregateCount.class);
+        String descrString = "QUERY-COUNT " + aggregateCategory.getName() + " WHERE 'a' = 'b'";
+        StatementDescriptor<AggregateCount> desc = new StatementDescriptor<>(aggregateCategory, descrString);
+        StatementDescriptorParser<AggregateCount> parser = new StatementDescriptorParser<>(storage, desc);
+        ParsedStatementImpl<AggregateCount> statement = (ParsedStatementImpl<AggregateCount>)parser.parse();
+        assertEquals(0, statement.getNumParams());
+        assertTrue(statement.getRawStatement() instanceof AggregateQuery);
+        Category<AggregateCount> capturedCategory = captor.getValue();
+        assertEquals(aggregateCategory, capturedCategory);
+        SuffixExpression expn = statement.getSuffixExpression();
+        assertNotNull(expn);
+        WhereExpression where = expn.getWhereExpn();
+        assertNotNull(where);
+        assertNull(expn.getSortExpn());
+        assertNull(expn.getLimitExpn());
+    }
+    
     @Test
     public void testParseQuerySimple() throws DescriptorParsingException {
         String descrString = "QUERY " + AgentInfoDAO.CATEGORY.getName();
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/storage/core/src/test/java/com/redhat/thermostat/storage/model/AggregateCountTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -0,0 +1,68 @@
+/*
+ * Copyright 2012, 2013 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;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import java.util.NoSuchElementException;
+
+import org.junit.Test;
+
+import com.redhat.thermostat.storage.core.Cursor;
+
+public class AggregateCountTest {
+
+    @Test
+    public void testCursor() {
+        AggregateCount c = new AggregateCount();
+        c.setCount(10);
+        Cursor<AggregateCount> cursor = c.getCursor();
+        assertTrue(cursor.hasNext());
+        AggregateCount actual = cursor.next();
+        assertEquals(10, actual.getCount());
+        assertFalse(cursor.hasNext());
+        try {
+            cursor.next();
+            fail("Should have thrown NoSuchElementException!");
+        } catch (NoSuchElementException e) {
+            // pass
+        }
+    }
+}
--- a/storage/mongo/src/main/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorage.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/mongo/src/main/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorage.java	Fri Aug 09 10:43:36 2013 +0200
@@ -55,6 +55,8 @@
 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.AggregateQuery;
+import com.redhat.thermostat.storage.core.AggregateQuery.AggregateFunction;
 import com.redhat.thermostat.storage.core.BackingStorage;
 import com.redhat.thermostat.storage.core.BasePut;
 import com.redhat.thermostat.storage.core.Category;
@@ -71,6 +73,8 @@
 import com.redhat.thermostat.storage.core.Replace;
 import com.redhat.thermostat.storage.core.StatementDescriptor;
 import com.redhat.thermostat.storage.core.Update;
+import com.redhat.thermostat.storage.model.AggregateCount;
+import com.redhat.thermostat.storage.model.AggregateResult;
 import com.redhat.thermostat.storage.model.Pojo;
 import com.redhat.thermostat.storage.query.Expression;
 
@@ -80,6 +84,21 @@
  * In this implementation, each CATEGORY is given a distinct collection.
  */
 public class MongoStorage implements BackingStorage {
+    
+    private class MongoCountQuery<T extends Pojo> extends AggregateQuery<T> {
+        
+        private final Category<T> category;
+        
+        private MongoCountQuery(MongoQuery<T> queryToAggregate, Category<T> category) {
+            super(AggregateFunction.COUNT, queryToAggregate);
+            this.category = category;
+        }
+
+        @Override
+        public Cursor<T> execute() {
+            return executeGetCount(category, (MongoQuery<T>)this.queryToAggregate);
+        }
+    }
 
     private class MongoAdd extends BasePut implements Add {
 
@@ -140,6 +159,12 @@
     private CountDownLatch connectedLatch;
     private UUID agentId;
 
+    // For testing only
+    MongoStorage(DB db, CountDownLatch latch) {
+        this.db = db;
+        this.connectedLatch = latch;
+    }
+    
     public MongoStorage(StartupConfiguration conf) {
         conn = new MongoConnection(conf);
         connectedLatch = new CountDownLatch(1);
@@ -160,6 +185,18 @@
         });
     }
 
+    public <T extends Pojo> Cursor<T> executeGetCount(Category<T> category, MongoQuery<T> queryToAggregate) {
+        DBCollection coll = getCachedCollection(category);
+        long count = 0L;
+        DBObject query = queryToAggregate.getGeneratedQuery();
+        if (coll != null) {
+            count = coll.getCount(query);
+        }
+        AggregateCount result = new AggregateCount();
+        result.setCount(count);
+        return result.getCursor();
+    }
+
     @Override
     public Connection getConnection() {
         return conn;
@@ -260,6 +297,11 @@
     
     @Override
     public void registerCategory(Category<?> category) {
+        Class<?> dataClass = category.getDataClass();
+        if (AggregateResult.class.isAssignableFrom(dataClass)) {
+            // adapted aggregate category, no need to actually register
+            return;
+        }
         String name = category.getName();
         if (collectionCache.containsKey(name)) {
             throw new IllegalStateException("Category may only be associated with one backend.");
@@ -366,5 +408,18 @@
         return PreparedStatementFactory.getInstance(this, statementDesc);
     }
 
+    @Override
+    public <T extends Pojo> Query<T> createAggregateQuery(
+            AggregateFunction function, Category<T> category) {
+        switch (function) {
+        case COUNT:
+            MongoQuery<T> query = (MongoQuery<T>)createQuery(category);
+            return new MongoCountQuery<>(query, category);
+        default:
+            throw new IllegalStateException("function not supported: "
+                    + function);
+        }
+    }
+
 }
 
--- a/storage/mongo/src/test/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorageTest.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/storage/mongo/src/test/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorageTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -47,9 +47,9 @@
 import static org.mockito.Matchers.eq;
 import static org.mockito.Matchers.same;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
-import static org.mockito.Mockito.times;
 
 import java.io.ByteArrayInputStream;
 import java.io.InputStream;
@@ -57,6 +57,7 @@
 import java.util.LinkedHashSet;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
 
 import org.junit.After;
 import org.junit.Before;
@@ -81,6 +82,7 @@
 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.CategoryAdapter;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Entity;
 import com.redhat.thermostat.storage.core.Key;
@@ -88,7 +90,10 @@
 import com.redhat.thermostat.storage.core.Put;
 import com.redhat.thermostat.storage.core.Query;
 import com.redhat.thermostat.storage.core.Update;
+import com.redhat.thermostat.storage.dao.HostInfoDAO;
+import com.redhat.thermostat.storage.model.AggregateCount;
 import com.redhat.thermostat.storage.model.BasePojo;
+import com.redhat.thermostat.storage.model.HostInfo;
 import com.redhat.thermostat.storage.query.Expression;
 import com.redhat.thermostat.storage.query.ExpressionFactory;
 
@@ -224,6 +229,18 @@
         emptyTestCollection = null;
         cursor = null;
     }
+    
+    @Test
+    public void testRegisterCategory() throws Exception {
+        DB db = PowerMockito.mock(DB.class);
+        CountDownLatch latch = new CountDownLatch(1);
+        MongoStorage storage = new MongoStorage(db, latch);
+        latch.countDown();
+        storage.registerCategory(HostInfoDAO.hostInfoCategory);
+        Category<AggregateCount> countCat = new CategoryAdapter<HostInfo, AggregateCount>(HostInfoDAO.hostInfoCategory).getAdapted(AggregateCount.class);
+        storage.registerCategory(countCat);
+        verify(db).collectionExists(eq(HostInfoDAO.hostInfoCategory.getName()));
+    }
 
     @Test
     public void verifyFindAllReturnsCursor() throws Exception {
--- a/web/client/src/main/java/com/redhat/thermostat/web/client/internal/WebStorage.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/web/client/src/main/java/com/redhat/thermostat/web/client/internal/WebStorage.java	Fri Aug 09 10:43:36 2013 +0200
@@ -516,11 +516,13 @@
     public void registerCategory(Category<?> category) throws StorageException {
         NameValuePair nameParam = new BasicNameValuePair("name",
                 category.getName());
+        NameValuePair dataClassParam = new BasicNameValuePair("data-class",
+                category.getDataClass().getName());
         
         NameValuePair categoryParam = new BasicNameValuePair("category",
                 gson.toJson(category));
         List<NameValuePair> formparams = Arrays
-                .asList(nameParam, categoryParam);
+                .asList(nameParam, categoryParam, dataClassParam);
         try (CloseableHttpEntity entity = post(endpoint + "/register-category",
                 formparams)) {
             Reader reader = getContentAsReader(entity);
--- a/web/common/src/test/java/com/redhat/thermostat/web/common/ThermostatGSONConverterTest.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/web/common/src/test/java/com/redhat/thermostat/web/common/ThermostatGSONConverterTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -44,6 +44,7 @@
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
 import com.redhat.thermostat.storage.model.AgentInformation;
+import com.redhat.thermostat.storage.model.AggregateCount;
 import com.redhat.thermostat.storage.model.Pojo;
 
 public class ThermostatGSONConverterTest {
@@ -89,4 +90,15 @@
         assertEquals("testing", actual[0].getAgentId());
         assertEquals(true, actual[0].isAlive());
     }
+    
+    @Test
+    public void canSerializeDeserializeAggregateCount() {
+        long expectedCount = 3333000333L;
+        AggregateCount count = new AggregateCount();
+        count.setCount(expectedCount);
+        String jsonStr = gson.toJson(count);
+        // now do the reverse
+        AggregateCount c2 = gson.fromJson(jsonStr, AggregateCount.class);
+        assertEquals(expectedCount, c2.getCount());
+    }
 }
--- a/web/server/src/main/java/com/redhat/thermostat/web/server/WebStorageEndPoint.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/web/server/src/main/java/com/redhat/thermostat/web/server/WebStorageEndPoint.java	Fri Aug 09 10:43:36 2013 +0200
@@ -47,6 +47,7 @@
 import java.security.Principal;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -74,7 +75,9 @@
 import com.redhat.thermostat.common.utils.LoggingUtils;
 import com.redhat.thermostat.shared.config.Configuration;
 import com.redhat.thermostat.shared.config.InvalidConfigurationException;
+import com.redhat.thermostat.storage.core.Categories;
 import com.redhat.thermostat.storage.core.Category;
+import com.redhat.thermostat.storage.core.CategoryAdapter;
 import com.redhat.thermostat.storage.core.Connection;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.DescriptorParsingException;
@@ -92,6 +95,7 @@
 import com.redhat.thermostat.storage.core.Update;
 import com.redhat.thermostat.storage.core.auth.DescriptorMetadata;
 import com.redhat.thermostat.storage.core.auth.StatementDescriptorMetadataFactory;
+import com.redhat.thermostat.storage.model.AggregateResult;
 import com.redhat.thermostat.storage.model.Pojo;
 import com.redhat.thermostat.storage.query.BinaryLogicalExpression;
 import com.redhat.thermostat.storage.query.BinaryLogicalOperator;
@@ -122,6 +126,7 @@
     private static final String TOKEN_MANAGER_TIMEOUT_PARAM = "token-manager-timeout";
     private static final String TOKEN_MANAGER_KEY = "token-manager";
     private static final String JETTY_JAAS_USER_PRINCIPAL_CLASS_NAME = "org.eclipse.jetty.plus.jaas.JAASUserPrincipal";
+    private static final String CATEGORY_KEY_FORMAT = "%s|%s";
 
     // our strings can contain non-ASCII characters. Use UTF-8
     // see also PR 1344
@@ -436,6 +441,7 @@
         
     }
 
+    @SuppressWarnings("unchecked") // need to adapt categories
     @WebStoragePathHandler( path = "register-category" )
     private synchronized void registerCategory(HttpServletRequest req, HttpServletResponse resp) throws IOException {
         if (! isAuthorized(req, resp, Roles.REGISTER_CATEGORY)) {
@@ -443,17 +449,45 @@
         }
         
         String categoryName = req.getParameter("name");
+        String dataClassName = req.getParameter("data-class");
+        // We need to index into the category map using name + data class since
+        // we have a different category for aggregate queries. For them the
+        // category name will be the same, but the data class will be different.
+        String categoryKey = String.format(CATEGORY_KEY_FORMAT, categoryName, dataClassName);
         String categoryParam = req.getParameter("category");
         int id;
-        if (categoryIds.containsKey(categoryName)) {
-            id = categoryIds.get(categoryName);
+        if (categoryIds.containsKey(categoryKey)) {
+            id = categoryIds.get(categoryKey);
         } else {
-            // The following has the side effect of registering the newly deserialized Category in the Categories class.
-            Category<?> category = gson.fromJson(categoryParam, Category.class);
-            storage.registerCategory(category);
-
+            Class<?> dataClass = getDataClassFromName(dataClassName);
+            Category<?> category = null;
+            if ((AggregateResult.class.isAssignableFrom(dataClass))) {
+                // Aggregate category case
+                Category<?> original = Categories.getByName(categoryName);
+                if (original == null) {
+                    // DAOs register categories when they are constructed. If we
+                    // end up triggering this we are in deep water. An aggregate
+                    // query was attempted before the underlying category is
+                    // registered at all? Not good!
+                    throw new IllegalStateException("Original category of aggregate not registered!");
+                }
+                // Adapt the original category to the one we want
+                @SuppressWarnings({ "rawtypes" })
+                CategoryAdapter adapter = new CategoryAdapter(original);
+                category = adapter.getAdapted(dataClass);
+                logger.log(Level.FINEST, "(id: " + currentCategoryId + ") not registering aggregate category " + category );
+            } else {
+                // Regular, non-aggregate category. Those categories we actually
+                // need to register with backing storage.
+                // 
+                // The following has the side effect of registering the newly
+                // deserialized Category in the Categories class.
+                category = gson.fromJson(categoryParam, Category.class);
+                storage.registerCategory(category);
+                logger.log(Level.FINEST, "(id: " + currentCategoryId + ") registered non-aggreate category: " + category);
+            }
             id = currentCategoryId;
-            categoryIds.put(categoryName, id);
+            categoryIds.put(categoryKey, id);
             categories.put(id, category);
             currentCategoryId++;
         }
@@ -464,6 +498,15 @@
         writer.flush();
     }
 
+    private Class<?> getDataClassFromName(String dataClassName) {
+        try {
+            Class<?> clazz = Class.forName(dataClassName);
+            return clazz;
+        } catch (ClassNotFoundException e) {
+            throw new IllegalStateException("Unknown data class!");
+        }
+    }
+
     @WebStoragePathHandler( path = "put-pojo" )
     private void putPojo(HttpServletRequest req, HttpServletResponse resp) {
         String insertParam = req.getParameter("insert");
--- a/web/server/src/test/java/com/redhat/thermostat/web/server/KnownDescriptorRegistryTest.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/web/server/src/test/java/com/redhat/thermostat/web/server/KnownDescriptorRegistryTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -81,7 +81,7 @@
         // storage-core registers 9 queries; this module has
         // only storage-core as maven dep which registers queries.
         // see DAOImplStatementDescriptorRegistration
-        assertEquals(10, trustedDescs.size());
+        assertEquals(13, trustedDescs.size());
     }
     
     @Test
--- a/web/server/src/test/java/com/redhat/thermostat/web/server/WebStorageEndpointTest.java	Fri Aug 02 17:49:01 2013 +0200
+++ b/web/server/src/test/java/com/redhat/thermostat/web/server/WebStorageEndpointTest.java	Fri Aug 09 10:43:36 2013 +0200
@@ -95,9 +95,12 @@
 import com.google.gson.GsonBuilder;
 import com.google.gson.reflect.TypeToken;
 import com.redhat.thermostat.storage.core.Add;
+import com.redhat.thermostat.storage.core.AggregateQuery;
+import com.redhat.thermostat.storage.core.AggregateQuery.AggregateFunction;
 import com.redhat.thermostat.storage.core.BackingStorage;
 import com.redhat.thermostat.storage.core.Categories;
 import com.redhat.thermostat.storage.core.Category;
+import com.redhat.thermostat.storage.core.CategoryAdapter;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Entity;
 import com.redhat.thermostat.storage.core.Key;
@@ -112,7 +115,10 @@
 import com.redhat.thermostat.storage.core.Update;
 import com.redhat.thermostat.storage.core.auth.DescriptorMetadata;
 import com.redhat.thermostat.storage.core.auth.StatementDescriptorRegistration;
+import com.redhat.thermostat.storage.dao.HostInfoDAO;
+import com.redhat.thermostat.storage.model.AggregateCount;
 import com.redhat.thermostat.storage.model.BasePojo;
+import com.redhat.thermostat.storage.model.HostInfo;
 import com.redhat.thermostat.storage.model.Pojo;
 import com.redhat.thermostat.storage.query.BinarySetMembershipExpression;
 import com.redhat.thermostat.storage.query.Expression;
@@ -585,6 +591,110 @@
         }
     }
     
+    @SuppressWarnings({ "unchecked", "rawtypes" })
+    @Test
+    public void authorizedPreparedAggregateQuery() throws Exception {
+        String strDescriptor = "QUERY-COUNT " + category.getName();
+        DescriptorMetadata metadata = new DescriptorMetadata();
+        setupTrustedStatementRegistry(strDescriptor, metadata);
+        
+        Set<BasicRole> roles = new HashSet<>();
+        roles.add(new RolePrincipal(Roles.REGISTER_CATEGORY));
+        roles.add(new RolePrincipal(Roles.PREPARE_STATEMENT));
+        roles.add(new RolePrincipal(Roles.READ));
+        roles.add(new RolePrincipal(Roles.ACCESS_REALM));
+        UserPrincipal testUser = new UserPrincipal("ignored1");
+        testUser.setRoles(roles);
+        
+        final LoginService loginService = new TestJAASLoginService(testUser); 
+        port = FreePortFinder.findFreePort(new TryPort() {
+            
+            @Override
+            public void tryPort(int port) throws Exception {
+                startServer(port, loginService);
+            }
+        });
+        
+        AggregateCount count = new AggregateCount();
+        count.setCount(500);
+        // prepare-statement does this under the hood
+        Query<AggregateCount> mockMongoQuery = mock(AggregateQuery.class);
+        Category<AggregateCount> adapted = new CategoryAdapter(category).getAdapted(AggregateCount.class);
+        registerCategory(adapted, "no-matter", "no-matter");
+        when(mockStorage.createAggregateQuery(eq(AggregateFunction.COUNT), eq(adapted))).thenReturn(mockMongoQuery);
+
+        Cursor<AggregateCount> cursor = mock(Cursor.class);
+        when(cursor.hasNext()).thenReturn(true).thenReturn(false);
+        when(cursor.next()).thenReturn(count);
+        
+        PreparedStatement mockPreparedQuery = mock(PreparedStatement.class);
+        when(mockStorage.prepareStatement(any(StatementDescriptor.class))).thenReturn(mockPreparedQuery);
+        
+        ParsedStatement mockParsedStatement = mock(ParsedStatement.class);
+        when(mockParsedStatement.getNumParams()).thenReturn(0);
+        when(mockParsedStatement.patchStatement(any(PreparedParameter[].class))).thenReturn(mockMongoQuery);
+        when(mockPreparedQuery.getParsedStatement()).thenReturn(mockParsedStatement);
+        
+        // The web layer
+        when(mockPreparedQuery.executeQuery()).thenReturn(cursor);
+        // And the mongo layer
+        when(mockMongoQuery.execute()).thenReturn(cursor);
+
+        String endpoint = getEndpoint();
+        URL url = new URL(endpoint + "/prepare-statement");
+        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
+        conn.setRequestMethod("POST");
+        sendAuthentication(conn, "no-matter", "no-matter");
+        conn.setRequestProperty("Content-Type", "application/x-www-form-urlencoded");
+        conn.setDoInput(true);
+        conn.setDoOutput(true);
+        Gson gson = new GsonBuilder()
+                .registerTypeHierarchyAdapter(WebQueryResponse.class, new WebQueryResponseSerializer<>())
+                .registerTypeAdapter(Pojo.class, new ThermostatGSONConverter())
+                .registerTypeAdapter(WebPreparedStatement.class, new WebPreparedStatementSerializer())
+                .registerTypeAdapter(PreparedParameter.class, new PreparedParameterSerializer()).create();
+        OutputStreamWriter out = new OutputStreamWriter(conn.getOutputStream());
+        String body = "query-descriptor=" + URLEncoder.encode(strDescriptor, "UTF-8") + "&category-id=" + categoryId;
+        out.write(body + "\n");
+        out.flush();
+
+        Reader in = new InputStreamReader(conn.getInputStream());
+        WebPreparedStatementResponse response = gson.fromJson(in, WebPreparedStatementResponse.class);
+        assertEquals(0, response.getNumFreeVariables());
+        assertEquals(0, response.getStatementId());
+        assertEquals("application/json; charset=UTF-8", conn.getContentType());
+        
+        
+        
+        // now execute the query we've just prepared
+        WebPreparedStatement<AggregateCount> stmt = new WebPreparedStatement<>(0, 0);
+        
+        url = new URL(endpoint + "/query-execute");
+        HttpURLConnection conn2 = (HttpURLConnection) url.openConnection();
+        conn2.setRequestMethod("POST");
+        sendAuthentication(conn2, "no-matter", "no-matter");
+        conn2.setRequestProperty("Content-Type", "application/x-www-form-urlencoded");
+        conn2.setDoInput(true);
+        conn2.setDoOutput(true);
+        
+        out = new OutputStreamWriter(conn2.getOutputStream());
+        body = "prepared-stmt=" + gson.toJson(stmt, WebPreparedStatement.class);
+        out.write(body + "\n");
+        out.flush();
+
+        in = new InputStreamReader(conn2.getInputStream());
+        Type typeToken = new TypeToken<WebQueryResponse<AggregateCount>>(){}.getType();
+        WebQueryResponse<AggregateCount> result = gson.fromJson(in, typeToken);
+        AggregateCount[] results = result.getResultList();
+        assertEquals(1, results.length);
+        assertEquals(500, results[0].getCount());
+
+        assertEquals("application/json; charset=UTF-8", conn2.getContentType());
+        verify(mockMongoQuery).execute();
+        verify(mockMongoQuery).getWhereExpression();
+        verifyNoMoreInteractions(mockMongoQuery);
+    }
+    
     private void setupTrustedStatementRegistry(String strDescriptor, DescriptorMetadata metadata) {
         Set<String> descs = new HashSet<>();
         descs.add(strDescriptor);
@@ -641,6 +751,56 @@
         
         assertEquals(failMessage, HttpServletResponse.SC_FORBIDDEN, conn.getResponseCode());
     }
+    
+    @Test
+    public void authorizedRegisterCategoryTest() throws Exception {
+        Set<BasicRole> roles = new HashSet<>();
+        roles.add(new RolePrincipal(Roles.REGISTER_CATEGORY));
+        roles.add(new RolePrincipal(Roles.ACCESS_REALM));
+        UserPrincipal testUser = new UserPrincipal("ignored1");
+        testUser.setRoles(roles);
+        
+        final LoginService loginService = new TestJAASLoginService(testUser); 
+        port = FreePortFinder.findFreePort(new TryPort() {
+            
+            @Override
+            public void tryPort(int port) throws Exception {
+                startServer(port, loginService);
+            }
+        });
+        Category<HostInfo> wantedCategory = HostInfoDAO.hostInfoCategory;
+        Category<AggregateCount> aggregate = new CategoryAdapter<HostInfo, AggregateCount>(wantedCategory).getAdapted(AggregateCount.class);
+        
+        // First the originating category has to be registered, then the adapted
+        // one.
+        Integer realId = registerCategoryAndGetId(wantedCategory, "no-matter", "no-matter");
+        Integer aggregateId = registerCategoryAndGetId(aggregate, "no-matter", "no-matter");
+        
+        assertTrue("Aggregate categories need their own ID", aggregateId != realId);
+        
+        verify(mockStorage).registerCategory(eq(wantedCategory));
+        verifyNoMoreInteractions(mockStorage);
+    }
+    
+    private Integer registerCategoryAndGetId(Category<?> cat, String username, String password) throws Exception {
+        String endpoint = getEndpoint();
+        URL url = new URL(endpoint + "/register-category");
+        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
+        conn.setRequestMethod("POST");
+        sendAuthentication(conn, username, password);
+
+        conn.setDoOutput(true);
+        conn.setDoInput(true);
+        conn.setRequestProperty("Content-Type", "application/x-www-form-urlencoded");
+        OutputStreamWriter out = new OutputStreamWriter(conn.getOutputStream());
+        Gson gson = new Gson();
+        out.write("name=" + cat.getName() + "&data-class=" + cat.getDataClass().getName() + "&category=" + gson.toJson(cat));
+        out.flush();
+        assertEquals(200, conn.getResponseCode());
+        Reader reader = new InputStreamReader(conn.getInputStream());
+        Integer id = gson.fromJson(reader, Integer.class);
+        return id;
+    }
 
     @Test
     public void authorizedReplacePutPojo() throws Exception {
@@ -1109,6 +1269,10 @@
     }
 
     private void registerCategory(String username, String password) {
+        registerCategory(category, username, password);
+    }
+    
+    private void registerCategory(Category<?> category, String username, String password) {
         try {
             String endpoint = getEndpoint();
             URL url = new URL(endpoint + "/register-category");
@@ -1126,6 +1290,8 @@
             writer.write(URLEncoder.encode(category.getName(), enc));
             writer.write("&category=");
             writer.write(URLEncoder.encode(gson.toJson(category), enc));
+            writer.write("&data-class=");
+            writer.write(URLEncoder.encode(category.getDataClass().getName(), enc));
             writer.flush();
 
             InputStream in = conn.getInputStream();