4 files changed, 30 insertions, 13 deletions

diff --git a/common/storage.h b/common/storage.h
index 589c738..71e9401 100644
--- a/common/storage.h
+++ b/common/storage.h
@@ -189,6 +189,9 @@ public:
     static QByteArray getTypeFromRevision(const Transaction &, qint64 revision);
     static void recordRevision(Transaction &, qint64 revision, const QByteArray &uid, const QByteArray &type);
     static void removeRevision(Transaction &, qint64 revision);
+    static void recordUid(DataStore::Transaction &transaction, const QByteArray &uid);
+    static void removeUid(DataStore::Transaction &transaction, const QByteArray &uid);
+    static void getUids(const Transaction &, const std::function<void(const QByteArray &uid)> &);
 
     bool exists() const;
 
diff --git a/common/storage/entitystore.cpp b/common/storage/entitystore.cpp
index bb7c33c..4afb407 100644
--- a/common/storage/entitystore.cpp
+++ b/common/storage/entitystore.cpp
@@ -204,6 +204,7 @@ bool EntityStore::add(const QByteArray &type, const ApplicationDomain::Applicati
             [&](const DataStore::Error &error) { SinkWarningCtx(d->logCtx) << "Failed to write entity" << entity.identifier() << newRevision; });
     DataStore::setMaxRevision(d->transaction, newRevision);
     DataStore::recordRevision(d->transaction, newRevision, entity.identifier(), type);
+    DataStore::recordUid(d->transaction, entity.identifier());
     SinkTraceCtx(d->logCtx) << "Wrote entity: " << entity.identifier() << type << newRevision;
     return true;
 }
@@ -327,6 +328,7 @@ bool EntityStore::remove(const QByteArray &type, const QByteArray &uid, bool rep
             [&](const DataStore::Error &error) { SinkWarningCtx(d->logCtx) << "Failed to write entity" << uid << newRevision; });
     DataStore::setMaxRevision(d->transaction, newRevision);
     DataStore::recordRevision(d->transaction, newRevision, uid, type);
+    DataStore::removeUid(d->transaction, uid);
     return true;
 }
 
@@ -581,19 +583,7 @@ ApplicationDomain::ApplicationDomainType EntityStore::readPrevious(const QByteAr
 
 void EntityStore::readAllUids(const QByteArray &type, const std::function<void(const QByteArray &uid)> callback)
 {
-    //TODO use a uid index instead
-    auto db = DataStore::mainDatabase(d->getTransaction(), type);
-    QByteArray lastUid;
-    db.scan("",
-        [&](const QByteArray &key, const QByteArray &) -> bool {
-            const auto uid = Sink::Storage::DataStore::uidFromKey(key);
-            if (uid != lastUid) {
-                lastUid = uid;
-                callback(uid);
-            }
-            return true;
-        },
-        [&](const Sink::Storage::DataStore::Error &error) { SinkWarningCtx(d->logCtx) << "Failed to read current value from storage: " << error.message; });
+    DataStore::getUids(d->getTransaction(), callback);
 }
 
 bool EntityStore::contains(const QByteArray &type, const QByteArray &uid)
diff --git a/common/storage_common.cpp b/common/storage_common.cpp
index 07d0f04..81a38c7 100644
--- a/common/storage_common.cpp
+++ b/common/storage_common.cpp
@@ -146,6 +146,24 @@ void DataStore::removeRevision(DataStore::Transaction &transaction, qint64 revis
     transaction.openDatabase("revisionType").remove(QByteArray::number(revision));
 }
 
+void DataStore::recordUid(DataStore::Transaction &transaction, const QByteArray &uid)
+{
+    transaction.openDatabase("uids").write(uid, "");
+}
+
+void DataStore::removeUid(DataStore::Transaction &transaction, const QByteArray &uid)
+{
+    transaction.openDatabase("uids").remove(uid);
+}
+
+void DataStore::getUids(const Transaction &transaction, const std::function<void(const QByteArray &uid)> &callback)
+{
+    transaction.openDatabase("uids").scan("", [&] (const QByteArray &key, const QByteArray &) {
+        callback(key);
+        return true;
+    });
+}
+
 bool DataStore::isInternalKey(const char *key)
 {
     return key && strncmp(key, s_internalPrefix, s_internalPrefixSize) == 0;
diff --git a/tests/entitystoretest.cpp b/tests/entitystoretest.cpp
index 0faf500..925ca7e 100644
--- a/tests/entitystoretest.cpp
+++ b/tests/entitystoretest.cpp
@@ -44,13 +44,19 @@ private slots:
         mail2.setExtractedMessageId("messageid2");
         mail2.setExtractedSubject("foo");
 
+        auto mail3 = ApplicationDomain::ApplicationDomainType::createEntity<ApplicationDomain::Mail>("res1");
+        mail3.setExtractedMessageId("messageid2");
+        mail3.setExtractedSubject("foo");
+
         store.startTransaction(Storage::DataStore::ReadWrite);
         store.add("mail", mail, false, [] (const ApplicationDomain::ApplicationDomainType &) {});
         store.add("mail", mail2, false, [] (const ApplicationDomain::ApplicationDomainType &) {});
+        store.add("mail", mail3, false, [] (const ApplicationDomain::ApplicationDomainType &) {});
 
         mail.setExtractedSubject("foo");
 
         store.modify("mail", mail, {}, false, [] (const ApplicationDomain::ApplicationDomainType &, ApplicationDomain::ApplicationDomainType &) {});
+        store.remove("mail", mail3.identifier(), false, [] (const ApplicationDomain::ApplicationDomainType &) {});
         store.commitTransaction();
 
         store.startTransaction(Storage::DataStore::ReadOnly);