4 files changed, 41 insertions, 12 deletions

diff --git a/common/storage.h b/common/storage.h
index 8c129df..c39b904 100644
--- a/common/storage.h
+++ b/common/storage.h
@@ -198,9 +198,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)> &);
+    static void recordUid(DataStore::Transaction &transaction, const QByteArray &uid, const QByteArray &type);
+    static void removeUid(DataStore::Transaction &transaction, const QByteArray &uid, const QByteArray &type);
+    static void getUids(const QByteArray &type, 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 1ac87b7..5514e31 100644
--- a/common/storage/entitystore.cpp
+++ b/common/storage/entitystore.cpp
@@ -267,7 +267,7 @@ bool EntityStore::add(const QByteArray &type, ApplicationDomain::ApplicationDoma
             [&](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());
+    DataStore::recordUid(d->transaction, entity.identifier(), type);
     SinkTraceCtx(d->logCtx) << "Wrote entity: " << entity.identifier() << type << newRevision;
     return true;
 }
@@ -379,7 +379,7 @@ bool EntityStore::remove(const QByteArray &type, const Sink::ApplicationDomain::
             [&](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);
+    DataStore::removeUid(d->transaction, uid, type);
     return true;
 }
 
@@ -638,7 +638,7 @@ ApplicationDomain::ApplicationDomainType EntityStore::readPrevious(const QByteAr
 
 void EntityStore::readAllUids(const QByteArray &type, const std::function<void(const QByteArray &uid)> callback)
 {
-    DataStore::getUids(d->getTransaction(), callback);
+    DataStore::getUids(type, 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 8603787..630dae9 100644
--- a/common/storage_common.cpp
+++ b/common/storage_common.cpp
@@ -156,19 +156,19 @@ void DataStore::removeRevision(DataStore::Transaction &transaction, qint64 revis
     transaction.openDatabase("revisionType").remove(QByteArray::number(revision));
 }
 
-void DataStore::recordUid(DataStore::Transaction &transaction, const QByteArray &uid)
+void DataStore::recordUid(DataStore::Transaction &transaction, const QByteArray &uid, const QByteArray &type)
 {
-    transaction.openDatabase("uids").write(uid, "");
+    transaction.openDatabase(type + "uids").write(uid, "");
 }
 
-void DataStore::removeUid(DataStore::Transaction &transaction, const QByteArray &uid)
+void DataStore::removeUid(DataStore::Transaction &transaction, const QByteArray &uid, const QByteArray &type)
 {
-    transaction.openDatabase("uids").remove(uid);
+    transaction.openDatabase(type + "uids").remove(uid);
 }
 
-void DataStore::getUids(const Transaction &transaction, const std::function<void(const QByteArray &uid)> &callback)
+void DataStore::getUids(const QByteArray &type, const Transaction &transaction, const std::function<void(const QByteArray &uid)> &callback)
 {
-    transaction.openDatabase("uids").scan("", [&] (const QByteArray &key, const QByteArray &) {
+    transaction.openDatabase(type + "uids").scan("", [&] (const QByteArray &key, const QByteArray &) {
         callback(key);
         return true;
     });
diff --git a/tests/storagetest.cpp b/tests/storagetest.cpp
index 9e9bad9..3368549 100644
--- a/tests/storagetest.cpp
+++ b/tests/storagetest.cpp
@@ -542,6 +542,35 @@ private slots:
         }
 
     }
+
+    void testRecordUid()
+    {
+        Sink::Storage::DataStore store(testDataPath, dbName, Sink::Storage::DataStore::ReadWrite);
+        auto transaction = store.createTransaction(Sink::Storage::DataStore::ReadWrite);
+        Sink::Storage::DataStore::recordUid(transaction, "uid1", "type");
+        Sink::Storage::DataStore::recordUid(transaction, "uid2", "type");
+        Sink::Storage::DataStore::recordUid(transaction, "uid3", "type2");
+
+        {
+            QVector<QByteArray> uids;
+            Sink::Storage::DataStore::getUids("type", transaction, [&](const QByteArray &r) {
+                uids << r;
+            });
+            QVector<QByteArray> expected{{"uid1"}, {"uid2"}};
+            QCOMPARE(uids, expected);
+        }
+
+        Sink::Storage::DataStore::removeUid(transaction, "uid2", "type");
+
+        {
+            QVector<QByteArray> uids;
+            Sink::Storage::DataStore::getUids("type", transaction, [&](const QByteArray &r) {
+                uids << r;
+            });
+            QVector<QByteArray> expected{{"uid1"}};
+            QCOMPARE(uids, expected);
+        }
+    }
 };
 
 QTEST_MAIN(StorageTest)