5 files changed, 244 insertions, 118 deletions

diff --git a/common/datastorequery.cpp b/common/datastorequery.cpp
index 41d962c..3ba8f40 100644
--- a/common/datastorequery.cpp
+++ b/common/datastorequery.cpp
@@ -72,7 +72,7 @@ class Source : public FilterBase {
         if (mIt == mIds.constEnd()) {
             return false;
         }
-        readEntity(*mIt, [callback](const Sink::ApplicationDomain::ApplicationDomainType &entity, Sink::Operation operation) {
+        readEntity(*mIt, [this, callback](const Sink::ApplicationDomain::ApplicationDomainType &entity, Sink::Operation operation) {
             SinkTraceCtx(mDatastore->mLogCtx) << "Source: Read entity: " << entity.identifier() << operationName(operation);
             callback({entity, operation});
         });
@@ -115,7 +115,7 @@ public:
     virtual bool next(const std::function<void(const ResultSet::Result &result)> &callback) Q_DECL_OVERRIDE {
         bool foundValue = false;
         while(!foundValue && mSource->next([this, callback, &foundValue](const ResultSet::Result &result) {
-                SinkTraceCtx(mDatastore->mLogCtx) << "Filter: " << result.entity.identifier() << result.operation;
+                SinkTraceCtx(mDatastore->mLogCtx) << "Filter: " << result.entity.identifier() << operationName(result.operation);
 
                 //Always accept removals. They can't match the filter since the data is gone.
                 if (result.operation == Sink::Operation_Removal) {
@@ -167,17 +167,11 @@ public:
 
         }
 
-        void process() {
-            if (operation == QueryBase::Reduce::Aggregator::Count) {
-                mResult = mResult.toInt() + 1;
-            } else {
-                Q_ASSERT(false);
-            }
-        }
-
         void process(const QVariant &value) {
             if (operation == QueryBase::Reduce::Aggregator::Collect) {
                 mResult = mResult.toList() << value;
+            } else if (operation == QueryBase::Reduce::Aggregator::Count) {
+                mResult = mResult.toInt() + 1;
             } else {
                 Q_ASSERT(false);
             }
@@ -196,8 +190,8 @@ public:
         QVariant mResult;
     };
 
-    QHash<QByteArray, QVariant> mAggregateValues;
     QSet<QByteArray> mReducedValues;
+    QHash<QByteArray, QByteArray> mSelectedValues;
     QByteArray mReductionProperty;
     QByteArray mSelectionProperty;
     QueryBase::Reduce::Selector::Comparator mSelectionComparator;
@@ -231,51 +225,80 @@ public:
         return false;
     }
 
+    QByteArray reduceOnValue(const QVariant &reductionValue, QMap<QByteArray, QVariant> &aggregateValues)
+    {
+        QVariant selectionResultValue;
+        QByteArray selectionResult;
+        auto results = indexLookup(mReductionProperty, reductionValue);
+        for (auto &aggregator : mAggregators) {
+            aggregator.reset();
+        }
+
+        for (const auto &r : results) {
+            readEntity(r, [&, this](const Sink::ApplicationDomain::ApplicationDomainType &entity, Sink::Operation operation) {
+                Q_ASSERT(operation != Sink::Operation_Removal);
+                for (auto &aggregator : mAggregators) {
+                    if (!aggregator.property.isEmpty()) {
+                        aggregator.process(entity.getProperty(aggregator.property));
+                    } else {
+                        aggregator.process(QVariant{});
+                    }
+                }
+                auto selectionValue = entity.getProperty(mSelectionProperty);
+                if (!selectionResultValue.isValid() || compare(selectionValue, selectionResultValue, mSelectionComparator)) {
+                    selectionResultValue = selectionValue;
+                    selectionResult = entity.identifier();
+                }
+            });
+        }
+
+        for (auto &aggregator : mAggregators) {
+            aggregateValues.insert(aggregator.resultProperty, aggregator.result());
+        }
+        return selectionResult;
+    }
+
     bool next(const std::function<void(const ResultSet::Result &)> &callback) Q_DECL_OVERRIDE {
         bool foundValue = false;
         while(!foundValue && mSource->next([this, callback, &foundValue](const ResultSet::Result &result) {
-                auto reductionValue = result.entity.getProperty(mReductionProperty);
                 if (result.operation == Sink::Operation_Removal) {
                     callback(result);
                     return false;
                 }
-                if (!mReducedValues.contains(getByteArray(reductionValue))) {
+                auto reductionValue = result.entity.getProperty(mReductionProperty);
+                const auto &reductionValueBa = getByteArray(reductionValue);
+                if (!mReducedValues.contains(reductionValueBa)) {
                     //Only reduce every value once.
-                    mReducedValues.insert(getByteArray(reductionValue));
-                    QVariant selectionResultValue;
-                    QByteArray selectionResult;
-                    auto results = indexLookup(mReductionProperty, reductionValue);
-                    for (auto &aggregator : mAggregators) {
-                        aggregator.reset();
-                    }
-
-                    QVariantList list;
-                    for (const auto &r : results) {
-                        readEntity(r, [&, this](const Sink::ApplicationDomain::ApplicationDomainType &entity, Sink::Operation operation) {
-                            for (auto &aggregator : mAggregators) {
-                                if (!aggregator.property.isEmpty()) {
-                                    aggregator.process(entity.getProperty(aggregator.property));
-                                } else {
-                                    aggregator.process();
-                                }
-                            }
-                            auto selectionValue = entity.getProperty(mSelectionProperty);
-                            if (!selectionResultValue.isValid() || compare(selectionValue, selectionResultValue, mSelectionComparator)) {
-                                selectionResultValue = selectionValue;
-                                selectionResult = entity.identifier();
-                            }
-                        });
-                    }
-
+                    mReducedValues.insert(reductionValueBa);
                     QMap<QByteArray, QVariant> aggregateValues;
-                    for (auto &aggregator : mAggregators) {
-                        aggregateValues.insert(aggregator.resultProperty, aggregator.result());
-                    }
+                    auto selectionResult = reduceOnValue(reductionValue, aggregateValues);
 
+                    mSelectedValues.insert(reductionValueBa, selectionResult);
                     readEntity(selectionResult, [&, this](const Sink::ApplicationDomain::ApplicationDomainType &entity, Sink::Operation operation) {
                         callback({entity, operation, aggregateValues});
                         foundValue = true;
                     });
+                } else {
+                    //During initial query, do nothing. The lookup above will take care of it.
+                    //During updates adjust the reduction according to the modification/addition or removal
+                    if (mIncremental) {
+                        //redo the reduction
+                        QMap<QByteArray, QVariant> aggregateValues;
+                        auto selectionResult = reduceOnValue(reductionValue, aggregateValues);
+
+                        //TODO if old and new are the same a modification would be enough
+                        auto oldSelectionResult = mSelectedValues.take(reductionValueBa);
+                        //remove old result
+                        readEntity(oldSelectionResult, [&, this](const Sink::ApplicationDomain::ApplicationDomainType &entity, Sink::Operation) {
+                            callback({entity, Sink::Operation_Removal});
+                        });
+
+                        //add new result
+                        mSelectedValues.insert(reductionValueBa, selectionResult);
+                        readEntity(selectionResult, [&, this](const Sink::ApplicationDomain::ApplicationDomainType &entity, Sink::Operation) {
+                            callback({entity, Sink::Operation_Creation, aggregateValues});
+                        });
+                    }
                 }
                 return false;
             }))
@@ -330,9 +353,23 @@ public:
 };
 
 DataStoreQuery::DataStoreQuery(const Sink::QueryBase &query, const QByteArray &type, EntityStore &store)
-    : mQuery(query), mType(type), mStore(store), mLogCtx(store.logContext().subContext("datastorequery"))
+    : mType(type), mStore(store), mLogCtx(store.logContext().subContext("datastorequery"))
+{
+    setupQuery(query);
+}
+
+DataStoreQuery::DataStoreQuery(const DataStoreQuery::State &state, const QByteArray &type, Sink::Storage::EntityStore &store)
+    : mType(type), mStore(store), mLogCtx(store.logContext().subContext("datastorequery"))
 {
-    setupQuery();
+    mCollector = state.mCollector;
+    mSource = state.mSource;
+
+    auto source = mCollector;
+    while (source) {
+        source->mDatastore = this;
+        source->mIncremental = true;
+        source = source->mSource;
+    }
 }
 
 DataStoreQuery::~DataStoreQuery()
@@ -340,6 +377,14 @@ DataStoreQuery::~DataStoreQuery()
 
 }
 
+DataStoreQuery::State::Ptr DataStoreQuery::getState()
+{
+    auto state = State::Ptr::create();
+    state->mSource = mSource;
+    state->mCollector = mCollector;
+    return state;
+}
+
 void DataStoreQuery::readEntity(const QByteArray &key, const BufferCallback &resultCallback)
 {
     mStore.readLatest(mType, key, resultCallback);
@@ -443,9 +488,10 @@ QByteArrayList DataStoreQuery::executeSubquery(const QueryBase &subquery)
     return ids;
 }
 
-void DataStoreQuery::setupQuery()
+void DataStoreQuery::setupQuery(const Sink::QueryBase &query_)
 {
-    auto baseFilters = mQuery.getBaseFilters();
+    auto query = query_;
+    auto baseFilters = query.getBaseFilters();
     for (const auto &k : baseFilters.keys()) {
         const auto comparator = baseFilters.value(k);
         if (comparator.value.canConvert<Query>()) {
@@ -454,44 +500,43 @@ void DataStoreQuery::setupQuery()
             baseFilters.insert(k, Query::Comparator(QVariant::fromValue(result), Query::Comparator::In));
         }
     }
-    mQuery.setBaseFilters(baseFilters);
+    query.setBaseFilters(baseFilters);
 
     FilterBase::Ptr baseSet;
-    QSet<QByteArray> remainingFilters = mQuery.getBaseFilters().keys().toSet();
+    QSet<QByteArray> remainingFilters = query.getBaseFilters().keys().toSet();
     QByteArray appliedSorting;
-    if (!mQuery.ids().isEmpty()) {
-        mSource = Source::Ptr::create(mQuery.ids().toVector(), this);
+    if (!query.ids().isEmpty()) {
+        mSource = Source::Ptr::create(query.ids().toVector(), this);
         baseSet = mSource;
     } else {
         QSet<QByteArray> appliedFilters;
 
-        auto resultSet = mStore.indexLookup(mType, mQuery, appliedFilters, appliedSorting);
+        auto resultSet = mStore.indexLookup(mType, query, appliedFilters, appliedSorting);
         remainingFilters = remainingFilters - appliedFilters;
 
         // We do a full scan if there were no indexes available to create the initial set.
         if (appliedFilters.isEmpty()) {
-            // TODO this should be replaced by an index lookup on the uid index
             mSource = Source::Ptr::create(mStore.fullScan(mType), this);
         } else {
             mSource = Source::Ptr::create(resultSet, this);
         }
         baseSet = mSource;
     }
-    if (!mQuery.getBaseFilters().isEmpty()) {
+    if (!query.getBaseFilters().isEmpty()) {
         auto filter = Filter::Ptr::create(baseSet, this);
         //For incremental queries the remaining filters are not sufficient
-        for (const auto &f : mQuery.getBaseFilters().keys()) {
-            filter->propertyFilter.insert(f, mQuery.getFilter(f));
+        for (const auto &f : query.getBaseFilters().keys()) {
+            filter->propertyFilter.insert(f, query.getFilter(f));
         }
         baseSet = filter;
     }
-    /* if (appliedSorting.isEmpty() && !mQuery.sortProperty.isEmpty()) { */
+    /* if (appliedSorting.isEmpty() && !query.sortProperty.isEmpty()) { */
     /*     //Apply manual sorting */
-    /*     baseSet = Sort::Ptr::create(baseSet, mQuery.sortProperty); */
+    /*     baseSet = Sort::Ptr::create(baseSet, query.sortProperty); */
     /* } */
 
     //Setup the rest of the filter stages on top of the base set
-    for (const auto &stage : mQuery.getFilterStages()) {
+    for (const auto &stage : query.getFilterStages()) {
         if (auto filter = stage.dynamicCast<Query::Filter>()) {
             auto f = Filter::Ptr::create(baseSet, this);
             f->propertyFilter = filter->propertyFilter;
@@ -521,7 +566,7 @@ QVector<QByteArray> DataStoreQuery::loadIncrementalResultSet(qint64 baseRevision
 
 ResultSet DataStoreQuery::update(qint64 baseRevision)
 {
-    SinkTraceCtx(mLogCtx) << "Executing query update to revision " << baseRevision;
+    SinkTraceCtx(mLogCtx) << "Executing query update from revision " << baseRevision;
     auto incrementalResultSet = loadIncrementalResultSet(baseRevision);
     SinkTraceCtx(mLogCtx) << "Incremental changes: " << incrementalResultSet;
     mSource->add(incrementalResultSet);
diff --git a/common/datastorequery.h b/common/datastorequery.h
index 5a47685..a797782 100644
--- a/common/datastorequery.h
+++ b/common/datastorequery.h
@@ -26,6 +26,7 @@
 
 class Source;
 class Bloom;
+class Reduce;
 class Filter;
 class FilterBase;
 
@@ -33,15 +34,25 @@ class DataStoreQuery {
     friend class FilterBase;
     friend class Source;
     friend class Bloom;
+    friend class Reduce;
     friend class Filter;
 public:
     typedef QSharedPointer<DataStoreQuery> Ptr;
 
+    struct State {
+        typedef QSharedPointer<State> Ptr;
+        QSharedPointer<FilterBase> mCollector;
+        QSharedPointer<Source> mSource;
+    };
+
     DataStoreQuery(const Sink::QueryBase &query, const QByteArray &type, Sink::Storage::EntityStore &store);
+    DataStoreQuery(const DataStoreQuery::State &state, const QByteArray &type, Sink::Storage::EntityStore &store);
     ~DataStoreQuery();
     ResultSet execute();
     ResultSet update(qint64 baseRevision);
 
+    State::Ptr getState();
+
 private:
 
     typedef std::function<bool(const Sink::ApplicationDomain::ApplicationDomainType &entity, Sink::Operation)> FilterFunction;
@@ -54,12 +65,10 @@ private:
     ResultSet createFilteredSet(ResultSet &resultSet, const FilterFunction &);
     QVector<QByteArray> loadIncrementalResultSet(qint64 baseRevision);
 
-    void setupQuery();
+    void setupQuery(const Sink::QueryBase &query_);
     QByteArrayList executeSubquery(const Sink::QueryBase &subquery);
 
-    Sink::QueryBase mQuery;
     const QByteArray mType;
-    bool mInitialQuery;
     QSharedPointer<FilterBase> mCollector;
     QSharedPointer<Source> mSource;
 
@@ -102,7 +111,8 @@ public:
     //Returns true for as long as a result is available
     virtual bool next(const std::function<void(const ResultSet::Result &)> &callback) = 0;
 
-    QSharedPointer<FilterBase> mSource;
+    FilterBase::Ptr mSource;
     DataStoreQuery *mDatastore;
+    bool mIncremental = false;
 };
 
diff --git a/common/queryrunner.cpp b/common/queryrunner.cpp
index 748320f..40880eb 100644
--- a/common/queryrunner.cpp
+++ b/common/queryrunner.cpp
@@ -33,6 +33,7 @@ struct ReplayResult {
     qint64 newRevision;
     qint64 replayedEntities;
     bool replayedAll;
+    DataStoreQuery::State::Ptr queryState;
 };
 
 /*
@@ -49,7 +50,7 @@ public:
     QueryWorker(const Sink::Query &query, const ResourceContext &context, const QByteArray &bufferType, const QueryRunnerBase::ResultTransformation &transformation, const Sink::Log::Context &logCtx);
     virtual ~QueryWorker();
 
-    ReplayResult executeIncrementalQuery(const Sink::Query &query, Sink::ResultProviderInterface<typename DomainType::Ptr> &resultProvider);
+    ReplayResult executeIncrementalQuery(const Sink::Query &query, Sink::ResultProviderInterface<typename DomainType::Ptr> &resultProvider, DataStoreQuery::State::Ptr state);
     ReplayResult executeInitialQuery(const Sink::Query &query, const typename DomainType::Ptr &parent, Sink::ResultProviderInterface<typename DomainType::Ptr> &resultProvider, int offset, int batchsize);
 
 private:
@@ -69,45 +70,41 @@ QueryRunner<DomainType>::QueryRunner(const Sink::Query &query, const Sink::Resou
         SinkWarningCtx(mLogCtx) << "A limited query without sorting is typically a bad idea, because there is no telling what you're going to get.";
     }
     auto guardPtr = QPointer<QObject>(&guard);
-    // We delegate loading of initial data to the result provider, so it can decide for itself what it needs to load.
-    mResultProvider->setFetcher([=](const typename DomainType::Ptr &parent) {
+    auto fetcher = [=](const typename DomainType::Ptr &parent) {
         const QByteArray parentId = parent ? parent->identifier() : QByteArray();
         SinkTraceCtx(mLogCtx) << "Running fetcher. Offset: " << mOffset[parentId] << " Batchsize: " << mBatchSize;
         auto resultProvider = mResultProvider;
-        if (query.synchronousQuery()) {
-            QueryWorker<DomainType> worker(query, mResourceContext, bufferType, mResultTransformation, mLogCtx);
-            const auto newRevisionAndReplayedEntities = worker.executeInitialQuery(query, parent, *resultProvider, mOffset[parentId], mBatchSize);
-            mOffset[parentId] += newRevisionAndReplayedEntities.replayedEntities;
-            resultProvider->setRevision(newRevisionAndReplayedEntities.newRevision);
-            resultProvider->initialResultSetComplete(parent, newRevisionAndReplayedEntities.replayedAll);
-        } else {
-            auto resultTransformation = mResultTransformation;
-            auto offset = mOffset[parentId];
-            auto batchSize = mBatchSize;
-            auto resourceContext = mResourceContext;
-            auto logCtx = mLogCtx;
-            //The lambda will be executed in a separate thread, so copy all arguments
-            async::run<ReplayResult>([resultTransformation, offset, batchSize, query, bufferType, resourceContext, resultProvider, parent, logCtx]() {
-                QueryWorker<DomainType> worker(query, resourceContext, bufferType, resultTransformation, logCtx);
-                const auto  newRevisionAndReplayedEntities = worker.executeInitialQuery(query, parent, *resultProvider, offset, batchSize);
-                return newRevisionAndReplayedEntities;
+        auto resultTransformation = mResultTransformation;
+        auto offset = mOffset[parentId];
+        auto batchSize = mBatchSize;
+        auto resourceContext = mResourceContext;
+        auto logCtx = mLogCtx;
+        const bool runAsync = !query.synchronousQuery();
+        //The lambda will be executed in a separate thread, so copy all arguments
+        async::run<ReplayResult>([resultTransformation, offset, batchSize, query, bufferType, resourceContext, resultProvider, parent, logCtx]() {
+            QueryWorker<DomainType> worker(query, resourceContext, bufferType, resultTransformation, logCtx);
+            return worker.executeInitialQuery(query, parent, *resultProvider, offset, batchSize);
+        }, runAsync)
+            .then([this, parentId, query, parent, resultProvider, guardPtr](const ReplayResult &result) {
+                if (!guardPtr) {
+                    qWarning() << "The parent object is already gone";
+                    return;
+                }
+                mInitialQueryComplete = true;
+                mQueryState = result.queryState;
+                mOffset[parentId] += result.replayedEntities;
+                // Only send the revision replayed information if we're connected to the resource, there's no need to start the resource otherwise.
+                if (query.liveQuery()) {
+                    mResourceAccess->sendRevisionReplayedCommand(result.newRevision);
+                }
+                resultProvider->setRevision(result.newRevision);
+                resultProvider->initialResultSetComplete(parent, result.replayedAll);
             })
-                .template then([this, parentId, query, parent, resultProvider, guardPtr](const ReplayResult &newRevisionAndReplayedEntities) {
-                    if (!guardPtr) {
-                        qWarning() << "The parent object is already gone";
-                        return;
-                    }
-                    mOffset[parentId] += newRevisionAndReplayedEntities.replayedEntities;
-                    // Only send the revision replayed information if we're connected to the resource, there's no need to start the resource otherwise.
-                    if (query.liveQuery()) {
-                        mResourceAccess->sendRevisionReplayedCommand(newRevisionAndReplayedEntities.newRevision);
-                    }
-                    resultProvider->setRevision(newRevisionAndReplayedEntities.newRevision);
-                    resultProvider->initialResultSetComplete(parent, newRevisionAndReplayedEntities.replayedAll);
-                })
-                .exec();
-        }
-    });
+            .exec();
+    };
+
+    // We delegate loading of initial data to the result provider, so it can decide for itself what it needs to load.
+    mResultProvider->setFetcher(fetcher);
 
     // In case of a live query we keep the runner for as long alive as the result provider exists
     if (query.liveQuery()) {
@@ -117,16 +114,26 @@ QueryRunner<DomainType>::QueryRunner(const Sink::Query &query, const Sink::Resou
             auto resultProvider = mResultProvider;
             auto resourceContext = mResourceContext;
             auto logCtx = mLogCtx;
-            return async::run<ReplayResult>([=]() {
+            auto state = mQueryState;
+            if (!mInitialQueryComplete) {
+                SinkWarningCtx(mLogCtx) << "Can't start the incremental query before the initial query is complete";
+                fetcher({});
+                return KAsync::null();
+            }
+            Q_ASSERT(!mQueryInProgress);
+            return KAsync::syncStart<void>([&] {
+                    mQueryInProgress = true;
+                })
+                .then(async::run<ReplayResult>([=]() {
                        QueryWorker<DomainType> worker(query, resourceContext, bufferType, mResultTransformation, logCtx);
-                       const auto newRevisionAndReplayedEntities = worker.executeIncrementalQuery(query, *resultProvider);
-                       return newRevisionAndReplayedEntities;
-                   })
-                .template then([query, this, resultProvider, guardPtr](const ReplayResult &newRevisionAndReplayedEntities) {
+                       return worker.executeIncrementalQuery(query, *resultProvider, state);
+                   }))
+                .then([query, this, resultProvider, guardPtr](const ReplayResult &newRevisionAndReplayedEntities) {
                     if (!guardPtr) {
                         qWarning() << "The parent object is already gone";
                         return;
                     }
+                    mQueryInProgress = false;
                     // Only send the revision replayed information if we're connected to the resource, there's no need to start the resource otherwise.
                     mResourceAccess->sendRevisionReplayedCommand(newRevisionAndReplayedEntities.newRevision);
                     resultProvider->setRevision(newRevisionAndReplayedEntities.newRevision);
@@ -201,14 +208,20 @@ void QueryWorker<DomainType>::resultProviderCallback(const Sink::Query &query, S
 }
 
 template <class DomainType>
-ReplayResult QueryWorker<DomainType>::executeIncrementalQuery(const Sink::Query &query, Sink::ResultProviderInterface<typename DomainType::Ptr> &resultProvider)
+ReplayResult QueryWorker<DomainType>::executeIncrementalQuery(const Sink::Query &query, Sink::ResultProviderInterface<typename DomainType::Ptr> &resultProvider, DataStoreQuery::State::Ptr state)
 {
     QTime time;
     time.start();
 
     const qint64 baseRevision = resultProvider.revision() + 1;
+    SinkTraceCtx(mLogCtx) << "Running query update from revision: " << baseRevision;
+
     auto entityStore = EntityStore{mResourceContext, mLogCtx};
-    auto preparedQuery = DataStoreQuery{query, ApplicationDomain::getTypeName<DomainType>(), entityStore};
+    if (!state) {
+        SinkWarningCtx(mLogCtx) << "No previous query state.";
+        return {0, 0, false, DataStoreQuery::State::Ptr{}};
+    }
+    auto preparedQuery = DataStoreQuery{*state, ApplicationDomain::getTypeName<DomainType>(), entityStore};
     auto resultSet = preparedQuery.update(baseRevision);
     SinkTraceCtx(mLogCtx) << "Filtered set retrieved. " << Log::TraceTime(time.elapsed());
     auto replayResult = resultSet.replaySet(0, 0, [this, query, &resultProvider](const ResultSet::Result &result) {
@@ -218,7 +231,7 @@ ReplayResult QueryWorker<DomainType>::executeIncrementalQuery(const Sink::Query
     SinkTraceCtx(mLogCtx) << "Replayed " << replayResult.replayedEntities << " results.\n"
         << (replayResult.replayedAll ? "Replayed all available results.\n" : "")
         << "Incremental query took: " << Log::TraceTime(time.elapsed());
-    return {entityStore.maxRevision(), replayResult.replayedEntities, replayResult.replayedAll};
+    return {entityStore.maxRevision(), replayResult.replayedEntities, replayResult.replayedAll, preparedQuery.getState()};
 }
 
 template <class DomainType>
@@ -251,14 +264,14 @@ ReplayResult QueryWorker<DomainType>::executeInitialQuery(
     SinkTraceCtx(mLogCtx) << "Replayed " << replayResult.replayedEntities << " results.\n"
         << (replayResult.replayedAll ? "Replayed all available results.\n" : "")
         << "Initial query took: " << Log::TraceTime(time.elapsed());
-    return {entityStore.maxRevision(), replayResult.replayedEntities, replayResult.replayedAll};
+
+    auto state = preparedQuery.getState();
+
+    return {entityStore.maxRevision(), replayResult.replayedEntities, replayResult.replayedAll, state};
 }
 
-template class QueryRunner<Sink::ApplicationDomain::Contact>;
-template class QueryRunner<Sink::ApplicationDomain::Folder>;
-template class QueryRunner<Sink::ApplicationDomain::Mail>;
-template class QueryRunner<Sink::ApplicationDomain::Event>;
-template class QueryWorker<Sink::ApplicationDomain::Contact>;
-template class QueryWorker<Sink::ApplicationDomain::Folder>;
-template class QueryWorker<Sink::ApplicationDomain::Mail>;
-template class QueryWorker<Sink::ApplicationDomain::Event>;
+#define REGISTER_TYPE(T) \
+    template class QueryRunner<T>; \
+    template class QueryWorker<T>; \
+
+SINK_REGISTER_TYPES()
diff --git a/common/queryrunner.h b/common/queryrunner.h
index 66dc68f..f5c7ead 100644
--- a/common/queryrunner.h
+++ b/common/queryrunner.h
@@ -24,6 +24,7 @@
 #include "resultprovider.h"
 #include "query.h"
 #include "log.h"
+#include "datastorequery.h"
 
 /**
  * Base clase because you can't have the Q_OBJECT macro in template classes
@@ -101,4 +102,7 @@ private:
     int mBatchSize;
     QObject guard;
     Sink::Log::Context mLogCtx;
+    DataStoreQuery::State::Ptr mQueryState;
+    bool mInitialQueryComplete = false;
+    bool mQueryInProgress = false;
 };
diff --git a/tests/querytest.cpp b/tests/querytest.cpp
index 0641c0d..328448f 100644
--- a/tests/querytest.cpp
+++ b/tests/querytest.cpp
@@ -731,7 +731,7 @@ private slots:
         Query query;
         query.setId("testLivequeryUnmatch");
         query.filter<Mail::Folder>(folder1);
-        query.reduce<Mail::ThreadId>(Query::Reduce::Selector::max<Mail::Date>()).count("count").collect<Mail::Sender>("senders");
+        query.reduce<Mail::ThreadId>(Query::Reduce::Selector::max<Mail::Date>()).count("count").collect<Mail::Folder>("folders");
         query.sort<Mail::Date>();
         query.setFlags(Query::LiveQuery);
         query.request<Mail::Unread>();
@@ -747,7 +747,61 @@ private slots:
         VERIFYEXEC(Sink::ResourceControl::flushMessageQueue("sink.dummy.instance1"));
         QTRY_COMPARE(model->rowCount(), 1);
         auto mail = model->data(model->index(0, 0, QModelIndex{}), Sink::Store::DomainObjectRole).value<Mail::Ptr>();
-        QCOMPARE(mail->getUnread(), true);
+        QTRY_COMPARE(mail->getUnread(), true);
+        QCOMPARE(mail->getProperty("count").toInt(), 1);
+        QCOMPARE(mail->getProperty("folders").toList().size(), 1);
+    }
+
+    void testReductionUpdate()
+    {
+        // Setup
+        auto folder1 = Folder::createEntity<Folder>("sink.dummy.instance1");
+        VERIFYEXEC(Sink::Store::create<Folder>(folder1));
+
+        auto folder2 = Folder::createEntity<Folder>("sink.dummy.instance1");
+        VERIFYEXEC(Sink::Store::create<Folder>(folder2));
+
+        QDateTime now{QDate{2017, 2, 3}, QTime{10, 0, 0}};
+        QDateTime later{QDate{2017, 2, 3}, QTime{11, 0, 0}};
+
+        auto mail1 = Mail::createEntity<Mail>("sink.dummy.instance1");
+        mail1.setExtractedMessageId("mail1");
+        mail1.setFolder(folder1);
+        mail1.setUnread(false);
+        mail1.setExtractedDate(now);
+        VERIFYEXEC(Sink::Store::create(mail1));
+
+        // Ensure all local data is processed
+        VERIFYEXEC(Sink::ResourceControl::flushMessageQueue("sink.dummy.instance1"));
+
+        Query query;
+        query.setId("testLivequeryUnmatch");
+        query.setFlags(Query::LiveQuery);
+        query.filter<Mail::Folder>(folder1);
+        query.reduce<Mail::Folder>(Query::Reduce::Selector::max<Mail::Date>()).count("count").collect<Mail::Folder>("folders");
+        query.sort<Mail::Date>();
+        query.request<Mail::Unread>();
+        query.request<Mail::MessageId>();
+
+        auto model = Sink::Store::loadModel<Mail>(query);
+        QTRY_COMPARE(model->rowCount(), 1);
+
+        //The leader should change to mail2 after the modification
+        {
+            auto mail2 = Mail::createEntity<Mail>("sink.dummy.instance1");
+            mail2.setExtractedMessageId("mail2");
+            mail2.setFolder(folder1);
+            mail2.setUnread(false);
+            mail2.setExtractedDate(later);
+            VERIFYEXEC(Sink::Store::create(mail2));
+        }
+
+        VERIFYEXEC(Sink::ResourceControl::flushMessageQueue("sink.dummy.instance1"));
+        QTRY_COMPARE(model->rowCount(), 1);
+        auto mail = model->data(model->index(0, 0, QModelIndex{}), Sink::Store::DomainObjectRole).value<Mail::Ptr>();
+        QTRY_COMPARE(mail->getMessageId(), QByteArray{"mail2"});
+        QCOMPARE(mail->getProperty("count").toInt(), 2);
+        QCOMPARE(mail->getProperty("folders").toList().size(), 2);
     }
 
     void testBloom()