11 files changed, 272 insertions, 51 deletions

diff --git a/common/changereplay.cpp b/common/changereplay.cpp
index 0adbd78..e81f52c 100644
--- a/common/changereplay.cpp
+++ b/common/changereplay.cpp
@@ -113,7 +113,9 @@ KAsync::Job<void> ChangeReplay::replayNextRevision()
                         if (uid.isEmpty() || type.isEmpty()) {
                             SinkErrorCtx(mLogCtx) << "Failed to read uid or type for revison: " << revision << uid << type;
                         } else {
-                            const auto key = DataStore::assembleKey(uid, revision);
+                            //const auto key = DataStore::assembleKey(uid, revision);
+                            // TODO: should not use internal representations
+                            const auto key = Storage::Key(Storage::Identifier::fromDisplayByteArray(uid), revision).toInternalByteArray();
                             QByteArray entityBuffer;
                             DataStore::mainDatabase(mMainStoreTransaction, type)
                                 .scan(key,
diff --git a/common/storage.h b/common/storage.h
index a8c486c..3cc5adf 100644
--- a/common/storage.h
+++ b/common/storage.h
@@ -22,8 +22,10 @@
 #pragma once
 
 #include "sink_export.h"
+#include "utils.h"
 #include <string>
 #include <functional>
+#include <QUuid>
 #include <QString>
 #include <QMap>
 
@@ -38,6 +40,151 @@ struct SINK_EXPORT DbLayout {
     Databases tables;
 };
 
+class Identifier
+{
+public:
+    // RFC 4122 Section 4.1.2 says 128 bits -> 16 bytes
+    static const constexpr size_t INTERNAL_REPR_SIZE = 16;
+    static const constexpr size_t DISPLAY_REPR_SIZE = 38;
+
+    Identifier() : uid(QUuid::createUuid()) {};
+
+    QByteArray toInternalByteArray() const
+    {
+        return uid.toRfc4122();
+    }
+
+    static Identifier fromInternalByteArray(const QByteArray &bytes)
+    {
+        Q_ASSERT(bytes.size() == INTERNAL_REPR_SIZE);
+        return Identifier(QUuid::fromRfc4122(bytes));
+    }
+
+    QString toDisplayString() const
+    {
+        return uid.toString();
+    }
+
+    QByteArray toDisplayByteArray() const
+    {
+        return uid.toByteArray();
+    }
+
+    static Identifier fromDisplayByteArray(const QByteArray &bytes)
+    {
+        Q_ASSERT(bytes.size() == DISPLAY_REPR_SIZE);
+        return Identifier(QUuid::fromString(QString::fromUtf8(bytes)));
+    }
+
+private:
+    explicit Identifier(const QUuid &uid) : uid(uid) {}
+    QUuid uid;
+};
+
+class Revision
+{
+public:
+    // qint64 has a 19 digit decimal representation
+    static const constexpr size_t INTERNAL_REPR_SIZE = 19;
+    static const constexpr size_t DISPLAY_REPR_SIZE = 19;
+
+    Revision(qint64 rev) : rev(rev) {}
+
+    QByteArray toInternalByteArray() const
+    {
+        return padNumber(rev);
+    }
+
+    static Revision fromInternalByteArray(const QByteArray &bytes)
+    {
+        Q_ASSERT(bytes.size() == INTERNAL_REPR_SIZE);
+        return Revision(bytes.toLongLong());
+    }
+
+    QString toDisplayString() const
+    {
+        return QString::fromUtf8(toInternalByteArray());
+    }
+
+    QByteArray toDisplayByteArray() const
+    {
+        return toInternalByteArray();
+    }
+
+    static Revision fromDisplayByteArray(const QByteArray &bytes)
+    {
+        Q_ASSERT(bytes.size() == DISPLAY_REPR_SIZE);
+        return fromInternalByteArray(bytes);
+    }
+
+    qint64 toQint64() const
+    {
+        return rev;
+    }
+
+private:
+    qint64 rev;
+};
+
+class Key
+{
+public:
+    static const constexpr size_t INTERNAL_REPR_SIZE = Identifier::INTERNAL_REPR_SIZE + Revision::INTERNAL_REPR_SIZE;
+    static const constexpr size_t DISPLAY_REPR_SIZE = Identifier::DISPLAY_REPR_SIZE + Revision::DISPLAY_REPR_SIZE;
+
+    Key(const Identifier &id, const Revision &rev) : id(id), rev(rev) {}
+
+    QByteArray toInternalByteArray() const
+    {
+        return id.toInternalByteArray() + rev.toInternalByteArray();
+    }
+
+    static Key fromInternalByteArray(const QByteArray &bytes)
+    {
+        Q_ASSERT(bytes.size() == INTERNAL_REPR_SIZE);
+        auto idBytes = bytes.mid(0, Identifier::INTERNAL_REPR_SIZE);
+        auto revBytes = bytes.mid(Identifier::INTERNAL_REPR_SIZE);
+        return Key(Identifier::fromInternalByteArray(idBytes), Revision::fromInternalByteArray(revBytes));
+    }
+
+    QString toDisplayString() const
+    {
+        return id.toDisplayString() + rev.toDisplayString();
+    }
+
+    QByteArray toDisplayByteArray() const
+    {
+        return id.toDisplayByteArray() + rev.toDisplayByteArray();
+    }
+
+    static Key fromDisplayByteArray(const QByteArray &bytes)
+    {
+        Q_ASSERT(bytes.size() == DISPLAY_REPR_SIZE);
+        auto idBytes = bytes.mid(0, Identifier::DISPLAY_REPR_SIZE);
+        auto revBytes = bytes.mid(Identifier::DISPLAY_REPR_SIZE);
+        return Key(Identifier::fromDisplayByteArray(idBytes), Revision::fromDisplayByteArray(revBytes));
+    }
+
+    const Identifier &identifier() const
+    {
+        return id;
+    }
+
+    const Revision &revision() const
+    {
+        return rev;
+    }
+
+    void setRevision(const Revision &newRev)
+    {
+        rev = newRev;
+    }
+
+private:
+    Identifier id;
+    Revision rev;
+};
+
 class SINK_EXPORT DataStore
 {
 public:
@@ -237,9 +384,9 @@ public:
     static bool isInternalKey(void *key, int keySize);
     static bool isInternalKey(const QByteArray &key);
 
-    static QByteArray assembleKey(const QByteArray &key, qint64 revision);
-    static QByteArray uidFromKey(const QByteArray &key);
-    static qint64 revisionFromKey(const QByteArray &key);
+    //static QByteArray assembleKey(const QByteArray &key, qint64 revision);
+    //static Identifier uidFromKey(const QByteArray &key);
+    //static qint64 revisionFromKey(const QByteArray &key);
 
     static NamedDatabase mainDatabase(const Transaction &, const QByteArray &type);
 
@@ -260,3 +407,6 @@ private:
 } // namespace Sink
 
 SINK_EXPORT QDebug& operator<<(QDebug &dbg, const Sink::Storage::DataStore::Error &error);
+SINK_EXPORT QDebug& operator<<(QDebug &dbg, const Sink::Storage::Identifier &);
+SINK_EXPORT QDebug& operator<<(QDebug &dbg, const Sink::Storage::Revision &);
+SINK_EXPORT QDebug& operator<<(QDebug &dbg, const Sink::Storage::Key &);
diff --git a/common/storage/entitystore.cpp b/common/storage/entitystore.cpp
index dd6bbf0..f74d3df 100644
--- a/common/storage/entitystore.cpp
+++ b/common/storage/entitystore.cpp
@@ -237,8 +237,10 @@ bool EntityStore::add(const QByteArray &type, ApplicationDomainType entity, bool
     flatbuffers::FlatBufferBuilder fbb;
     d->resourceContext.adaptorFactory(type).createBuffer(entity, fbb, metadataFbb.GetBufferPointer(), metadataFbb.GetSize());
 
+    const auto key = Key(Identifier::fromDisplayByteArray(entity.identifier()), newRevision);
+
     DataStore::mainDatabase(d->transaction, type)
-        .write(DataStore::assembleKey(entity.identifier(), newRevision), BufferUtils::extractBuffer(fbb),
+        .write(key.toInternalByteArray(), BufferUtils::extractBuffer(fbb),
             [&](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);
@@ -311,8 +313,10 @@ bool EntityStore::modify(const QByteArray &type, const ApplicationDomainType &cu
     flatbuffers::FlatBufferBuilder fbb;
     d->resourceContext.adaptorFactory(type).createBuffer(newEntity, fbb, metadataFbb.GetBufferPointer(), metadataFbb.GetSize());
 
+    const auto key = Key(Identifier::fromDisplayByteArray(newEntity.identifier()), newRevision);
+
     DataStore::mainDatabase(d->transaction, type)
-        .write(DataStore::assembleKey(newEntity.identifier(), newRevision), BufferUtils::extractBuffer(fbb),
+        .write(key.toInternalByteArray(), BufferUtils::extractBuffer(fbb),
             [&](const DataStore::Error &error) { SinkWarningCtx(d->logCtx) << "Failed to write entity" << newEntity.identifier() << newRevision; });
     DataStore::setMaxRevision(d->transaction, newRevision);
     DataStore::recordRevision(d->transaction, newRevision, newEntity.identifier(), type);
@@ -346,8 +350,10 @@ bool EntityStore::remove(const QByteArray &type, const ApplicationDomainType &cu
     flatbuffers::FlatBufferBuilder fbb;
     EntityBuffer::assembleEntityBuffer(fbb, metadataFbb.GetBufferPointer(), metadataFbb.GetSize(), 0, 0, 0, 0);
 
+    const auto key = Key(Identifier::fromDisplayByteArray(uid), newRevision);
+
     DataStore::mainDatabase(d->transaction, type)
-        .write(DataStore::assembleKey(uid, newRevision), BufferUtils::extractBuffer(fbb),
+        .write(key.toInternalByteArray(), BufferUtils::extractBuffer(fbb),
             [&](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);
@@ -428,7 +434,7 @@ QVector<QByteArray> EntityStore::fullScan(const QByteArray &type)
     DataStore::mainDatabase(d->getTransaction(), type)
         .scan(QByteArray(),
             [&](const QByteArray &key, const QByteArray &value) -> bool {
-                const auto uid = DataStore::uidFromKey(key);
+                const auto uid = Sink::Storage::Key::fromInternalByteArray(key).identifier().toDisplayByteArray();
                 if (keys.contains(uid)) {
                     //Not something that should persist if the replay works, so we keep a message for now.
                     SinkTraceCtx(d->logCtx) << "Multiple revisions for key: " << key;
@@ -479,16 +485,24 @@ void EntityStore::indexLookup(const QByteArray &type, const QByteArray &property
     /* }); */
 }
 
-void EntityStore::readLatest(const QByteArray &type, const QByteArray &uid, const std::function<void(const QByteArray &uid, const EntityBuffer &entity)> callback)
+void EntityStore::readLatest(const QByteArray &type, const QByteArray &key, const std::function<void(const QByteArray &uid, const EntityBuffer &entity)> callback)
 {
     Q_ASSERT(d);
-    Q_ASSERT(!uid.isEmpty());
+    Q_ASSERT(!key.isEmpty());
+    const auto internalKey = [&key]() {
+        if(key.size() == Identifier::DISPLAY_REPR_SIZE) {
+            return Identifier::fromDisplayByteArray(key).toInternalByteArray();
+        } else {
+            return Key::fromDisplayByteArray(key).toInternalByteArray();
+        }
+    }();
     auto db = DataStore::mainDatabase(d->getTransaction(), type);
-    db.findLatest(uid,
+    db.findLatest(internalKey,
         [=](const QByteArray &key, const QByteArray &value) {
-            callback(DataStore::uidFromKey(key), Sink::EntityBuffer(value.data(), value.size()));
+            const auto uid = Sink::Storage::Key::fromInternalByteArray(key).identifier().toDisplayByteArray();
+            callback(uid, Sink::EntityBuffer(value.data(), value.size()));
         },
-        [&](const DataStore::Error &error) { SinkWarningCtx(d->logCtx) << "Error during readLatest query: " << error.message << uid; });
+        [&](const DataStore::Error &error) { SinkWarningCtx(d->logCtx) << "Error during readLatest query: " << error.message << key; });
 }
 
 void EntityStore::readLatest(const QByteArray &type, const QByteArray &uid, const std::function<void(const ApplicationDomainType &)> callback)
@@ -516,12 +530,15 @@ ApplicationDomain::ApplicationDomainType EntityStore::readLatest(const QByteArra
     return dt;
 }
 
-void EntityStore::readEntity(const QByteArray &type, const QByteArray &key, const std::function<void(const QByteArray &uid, const EntityBuffer &entity)> callback)
+// TODO: check every usage
+void EntityStore::readEntity(const QByteArray &type, const QByteArray &displayKey, const std::function<void(const QByteArray &uid, const EntityBuffer &entity)> callback)
 {
+    const auto key = Key::fromDisplayByteArray(displayKey);
     auto db = DataStore::mainDatabase(d->getTransaction(), type);
-    db.scan(key,
+    db.scan(key.toInternalByteArray(),
         [=](const QByteArray &key, const QByteArray &value) -> bool {
-            callback(DataStore::uidFromKey(key), Sink::EntityBuffer(value.data(), value.size()));
+            const auto uid = Sink::Storage::Key::fromInternalByteArray(key).identifier().toDisplayByteArray();
+            callback(uid, Sink::EntityBuffer(value.data(), value.size()));
             return false;
         },
         [&](const DataStore::Error &error) { SinkWarningCtx(d->logCtx) << "Error during readEntity query: " << error.message << key; });
@@ -567,9 +584,10 @@ void EntityStore::readRevisions(qint64 baseRevision, const QByteArray &expectedT
             revisionCounter++;
             continue;
         }
-        const auto key = DataStore::assembleKey(uid, revisionCounter);
+        const auto key = Key(Identifier::fromDisplayByteArray(uid), revisionCounter);
+
         revisionCounter++;
-        callback(key);
+        callback(key.toDisplayByteArray());
     }
 }
 
@@ -577,16 +595,19 @@ void EntityStore::readPrevious(const QByteArray &type, const QByteArray &uid, qi
 {
     auto db = DataStore::mainDatabase(d->getTransaction(), type);
     qint64 latestRevision = 0;
-    db.scan(uid,
+    const auto internalUid = Identifier::fromDisplayByteArray(uid).toInternalByteArray();
+    db.scan(internalUid,
         [&latestRevision, revision](const QByteArray &key, const QByteArray &) -> bool {
-            const auto foundRevision = DataStore::revisionFromKey(key);
+            //const auto foundRevision = DataStore::revisionFromKey(key);
+            const auto foundRevision = Key::fromInternalByteArray(key).revision().toQint64();
             if (foundRevision < revision && foundRevision > latestRevision) {
                 latestRevision = foundRevision;
             }
             return true;
         },
         [&](const DataStore::Error &error) { SinkWarningCtx(d->logCtx) << "Failed to read current value from storage: " << error.message; }, true);
-    readEntity(type, DataStore::assembleKey(uid, latestRevision), callback);
+    const auto key = Key(Identifier::fromDisplayByteArray(uid), latestRevision);
+    readEntity(type, key.toDisplayByteArray(), callback);
 }
 
 void EntityStore::readPrevious(const QByteArray &type, const QByteArray &uid, qint64 revision, const std::function<void(const ApplicationDomainType &)> callback)
@@ -612,15 +633,22 @@ void EntityStore::readAllUids(const QByteArray &type, const std::function<void(c
 
 bool EntityStore::contains(const QByteArray &type, const QByteArray &uid)
 {
-    return DataStore::mainDatabase(d->getTransaction(), type).contains(uid);
+    // Because of pipeline using this function with new entities
+    // TODO: maybe modify pipeline instead?
+    if(uid.isEmpty()) {
+        return false;
+    }
+    const auto internalUid = Identifier::fromDisplayByteArray(uid).toInternalByteArray();
+    return DataStore::mainDatabase(d->getTransaction(), type).contains(internalUid);
 }
 
 bool EntityStore::exists(const QByteArray &type, const QByteArray &uid)
 {
     bool found = false;
     bool alreadyRemoved = false;
+    const auto internalUid = Identifier::fromDisplayByteArray(uid).toInternalByteArray();
     DataStore::mainDatabase(d->transaction, type)
-        .findLatest(uid,
+        .findLatest(internalUid,
             [&found, &alreadyRemoved](const QByteArray &key, const QByteArray &data) {
                 auto entity = GetEntity(data.data());
                 if (entity && entity->metadata()) {
diff --git a/common/storage_common.cpp b/common/storage_common.cpp
index 057dce4..b36ffcb 100644
--- a/common/storage_common.cpp
+++ b/common/storage_common.cpp
@@ -24,18 +24,40 @@
 #include "log.h"
 #include "utils.h"
 
+#include <QUuid>
+
 QDebug& operator<<(QDebug &dbg, const Sink::Storage::DataStore::Error &error)
 {
     dbg << error.message << "Code: " << error.code << "Db: " << error.store;
     return dbg;
 }
 
+QDebug& operator<<(QDebug &dbg, const Sink::Storage::Identifier &id)
+{
+    dbg << id.toDisplayString();
+    return dbg;
+}
+
+QDebug& operator<<(QDebug &dbg, const Sink::Storage::Revision &rev)
+{
+    dbg << rev.toDisplayString();
+    return dbg;
+}
+
+QDebug& operator<<(QDebug &dbg, const Sink::Storage::Key &key)
+{
+    dbg << key.toDisplayString();
+    return dbg;
+}
+
 namespace Sink {
 namespace Storage {
 
 static const char *s_internalPrefix = "__internal";
 static const int s_internalPrefixSize = strlen(s_internalPrefix);
 static const int s_lengthOfUid = 38;
+// RFC 4122 Section 4.1.2 says 128 bits -> 16 bytes
+//static const int s_lengthOfUid = 16;
 
 DbLayout::DbLayout()
 {
@@ -194,6 +216,7 @@ bool DataStore::isInternalKey(const QByteArray &key)
     return key.startsWith(s_internalPrefix);
 }
 
+/*
 QByteArray DataStore::assembleKey(const QByteArray &key, qint64 revision)
 {
     Q_ASSERT(revision <= 9223372036854775807);
@@ -201,15 +224,17 @@ QByteArray DataStore::assembleKey(const QByteArray &key, qint64 revision)
     return key + QByteArray::number(revision).rightJustified(19, '0', false);
 }
 
-QByteArray DataStore::uidFromKey(const QByteArray &key)
+//QByteArray DataStore::uidFromKey(const QByteArray &key)
+Identifier DataStore::uidFromKey(const QByteArray &key)
 {
-    return key.mid(0, s_lengthOfUid);
+    return Identifier::fromByteArray(key.mid(0, s_lengthOfUid));
 }
 
 qint64 DataStore::revisionFromKey(const QByteArray &key)
 {
     return key.mid(s_lengthOfUid + 1).toLongLong();
 }
+*/
 
 QByteArray DataStore::generateUid()
 {
diff --git a/common/synchronizer.cpp b/common/synchronizer.cpp
index 41ab1e9..f19c08a 100644
--- a/common/synchronizer.cpp
+++ b/common/synchronizer.cpp
@@ -616,7 +616,8 @@ KAsync::Job<void> Synchronizer::replay(const QByteArray &type, const QByteArray
     Q_ASSERT(mEntityStore->hasTransaction());
 
     const auto operation = metadataBuffer ? metadataBuffer->operation() : Sink::Operation_Creation;
-    const auto uid = Sink::Storage::DataStore::uidFromKey(key);
+    //const auto uid = Sink::Storage::DataStore::uidFromKey(key);
+    const auto uid = Sink::Storage::Key::fromDisplayByteArray(key).identifier().toDisplayByteArray();
     const auto modifiedProperties = metadataBuffer->modifiedProperties() ? BufferUtils::fromVector(*metadataBuffer->modifiedProperties()) : QByteArrayList();
     QByteArray oldRemoteId;
 
diff --git a/common/typeindex.cpp b/common/typeindex.cpp
index b18791f..0b78d59 100644
--- a/common/typeindex.cpp
+++ b/common/typeindex.cpp
@@ -25,8 +25,6 @@
 #include <QDateTime>
 #include <QDataStream>
 
-#include <cmath>
-
 using namespace Sink;
 
 static QByteArray getByteArray(const QVariant &value)
@@ -53,13 +51,6 @@ static QByteArray getByteArray(const QVariant &value)
     return "toplevel";
 }
 
-template <typename T>
-static QByteArray padNumber(T number)
-{
-    static T uint_num_digits = (T)std::log10(std::numeric_limits<T>::max()) + 1;
-    return QByteArray::number(number).rightJustified(uint_num_digits, '0');
-}
-
 static QByteArray toSortableByteArrayImpl(const QDateTime &date)
 {
     // Sort invalid last
diff --git a/common/utils.cpp b/common/utils.cpp
index 3c54db4..e98365c 100644
--- a/common/utils.cpp
+++ b/common/utils.cpp
@@ -22,4 +22,5 @@
 QByteArray Sink::createUuid()
 {
     return QUuid::createUuid().toByteArray();
+    //return QUuid::createUuid().toRfc4122();
 }
diff --git a/common/utils.h b/common/utils.h
index 253de61..7066d79 100644
--- a/common/utils.h
+++ b/common/utils.h
@@ -20,6 +20,17 @@
 
 #include <QByteArray>
 
+#include <cmath>
+
 namespace Sink {
-    QByteArray createUuid();
+
+QByteArray createUuid();
+
+template <typename T>
+static QByteArray padNumber(T number)
+{
+    static T uint_num_digits = (T)std::log10(std::numeric_limits<T>::max()) + 1;
+    return QByteArray::number(number).rightJustified(uint_num_digits, '0');
 }
+
+} // namespace Sink
diff --git a/sinksh/syntax_modules/sink_inspect.cpp b/sinksh/syntax_modules/sink_inspect.cpp
index 1a964cf..b496fe6 100644
--- a/sinksh/syntax_modules/sink_inspect.cpp
+++ b/sinksh/syntax_modules/sink_inspect.cpp
@@ -90,7 +90,8 @@ bool inspect(const QStringList &args, State &state)
 
         QSet<QByteArray> uids;
         db.scan("", [&] (const QByteArray &key, const QByteArray &data) {
-                    uids.insert(Sink::Storage::DataStore::uidFromKey(key));
+                    //uids.insert(Sink::Storage::DataStore::uidFromKey(key));
+                    uids.insert(Sink::Storage::Key::fromInternalByteArray(key).identifier().toDisplayByteArray());
                     return true;
                 },
                 [&](const Sink::Storage::DataStore::Error &e) {
diff --git a/tests/pipelinetest.cpp b/tests/pipelinetest.cpp
index 45e2fbb..7431deb 100644
--- a/tests/pipelinetest.cpp
+++ b/tests/pipelinetest.cpp
@@ -250,8 +250,9 @@ private slots:
         // Get uid of written entity
         auto keys = getKeys(instanceIdentifier(), "event.main");
         QCOMPARE(keys.size(), 1);
-        const auto key = keys.first();
-        const auto uid = Sink::Storage::DataStore::uidFromKey(key);
+        auto key = Sink::Storage::Key::fromInternalByteArray(keys.first());
+        //const auto uid = Sink::Storage::DataStore::uidFromKey(key);
+        const auto uid = key.identifier().toDisplayByteArray();
 
         // Execute the modification
         entityFbb.Clear();
@@ -260,8 +261,10 @@ private slots:
         pipeline.modifiedEntity(modifyCommand.constData(), modifyCommand.size());
         pipeline.commit();
 
+        key.setRevision(2);
+
         // Ensure we've got the new revision with the modification
-        auto buffer = getEntity(instanceIdentifier(), "event.main", Sink::Storage::DataStore::assembleKey(uid, 2));
+        auto buffer = getEntity(instanceIdentifier(), "event.main", key.toInternalByteArray());
         QVERIFY(!buffer.isEmpty());
         Sink::EntityBuffer entityBuffer(buffer.data(), buffer.size());
         auto adaptor = adaptorFactory->createAdaptor(entityBuffer.entity());
@@ -296,7 +299,8 @@ private slots:
         // Get uid of written entity
         auto keys = getKeys(instanceIdentifier(), "event.main");
         QCOMPARE(keys.size(), 1);
-        const auto uid = Sink::Storage::DataStore::uidFromKey(keys.first());
+        auto key = Sink::Storage::Key::fromInternalByteArray(keys.first());
+        const auto uid = key.identifier().toDisplayByteArray();
 
 
         // Create another operation inbetween
@@ -315,8 +319,10 @@ private slots:
         pipeline.modifiedEntity(modifyCommand.constData(), modifyCommand.size());
         pipeline.commit();
 
+        key.setRevision(3);
+
         // Ensure we've got the new revision with the modification
-        auto buffer = getEntity(instanceIdentifier(), "event.main", Sink::Storage::DataStore::assembleKey(uid, 3));
+        auto buffer = getEntity(instanceIdentifier(), "event.main", key.toInternalByteArray());
         QVERIFY(!buffer.isEmpty());
         Sink::EntityBuffer entityBuffer(buffer.data(), buffer.size());
         auto adaptor = adaptorFactory->createAdaptor(entityBuffer.entity());
@@ -337,7 +343,7 @@ private slots:
         auto result = getKeys(instanceIdentifier(), "event.main");
         QCOMPARE(result.size(), 1);
 
-        const auto uid = Sink::Storage::DataStore::uidFromKey(result.first());
+        const auto uid = Sink::Storage::Key::fromInternalByteArray(result.first()).identifier().toDisplayByteArray();
 
         // Delete entity
         auto deleteCommand = deleteEntityCommand(uid, 1);
@@ -373,21 +379,23 @@ private slots:
             QCOMPARE(testProcessor->newUids.size(), 1);
             QCOMPARE(testProcessor->newRevisions.size(), 1);
             // Key doesn't contain revision and is just the uid
-            QCOMPARE(testProcessor->newUids.at(0), Sink::Storage::DataStore::uidFromKey(testProcessor->newUids.at(0)));
+            const auto uid = Sink::Storage::Key::fromDisplayByteArray(testProcessor->newUids.at(0)).identifier().toDisplayByteArray();
+            QCOMPARE(testProcessor->newUids.at(0), uid);
         }
         pipeline.commit();
         entityFbb.Clear();
         pipeline.startTransaction();
         auto keys = getKeys(instanceIdentifier(), "event.main");
         QCOMPARE(keys.size(), 1);
-        const auto uid = Sink::Storage::DataStore::uidFromKey(keys.first());
+        const auto uid = Sink::Storage::Key::fromDisplayByteArray(keys.first()).identifier().toDisplayByteArray();
         {
             auto modifyCommand = modifyEntityCommand(createEvent(entityFbb, "summary2"), uid, 1);
             pipeline.modifiedEntity(modifyCommand.constData(), modifyCommand.size());
             QCOMPARE(testProcessor->modifiedUids.size(), 1);
             QCOMPARE(testProcessor->modifiedRevisions.size(), 1);
             // Key doesn't contain revision and is just the uid
-            QCOMPARE(testProcessor->modifiedUids.at(0), Sink::Storage::DataStore::uidFromKey(testProcessor->modifiedUids.at(0)));
+            const auto uid2 = Sink::Storage::Key::fromDisplayByteArray(testProcessor->modifiedUids.at(0)).identifier().toDisplayByteArray();
+            QCOMPARE(testProcessor->modifiedUids.at(0), uid2);
         }
         pipeline.commit();
         entityFbb.Clear();
@@ -399,7 +407,8 @@ private slots:
             QCOMPARE(testProcessor->deletedUids.size(), 1);
             QCOMPARE(testProcessor->deletedSummaries.size(), 1);
             // Key doesn't contain revision and is just the uid
-            QCOMPARE(testProcessor->deletedUids.at(0), Sink::Storage::DataStore::uidFromKey(testProcessor->deletedUids.at(0)));
+            const auto uid2 = Sink::Storage::Key::fromDisplayByteArray(testProcessor->modifiedUids.at(0)).identifier().toDisplayByteArray();
+            QCOMPARE(testProcessor->deletedUids.at(0), uid2);
             QCOMPARE(testProcessor->deletedSummaries.at(0), QByteArray("summary2"));
         }
     }
diff --git a/tests/storagetest.cpp b/tests/storagetest.cpp
index bca91b1..40492f0 100644
--- a/tests/storagetest.cpp
+++ b/tests/storagetest.cpp
@@ -456,8 +456,10 @@ private slots:
         auto db = transaction.openDatabase("test", nullptr, false);
         const auto uid = "{c5d06a9f-1534-4c52-b8ea-415db68bdadf}";
         //Ensure we can sort 1 and 10 properly (by default string comparison 10 comes before 6)
-        db.write(Sink::Storage::DataStore::assembleKey(uid, 6), "value1");
-        db.write(Sink::Storage::DataStore::assembleKey(uid, 10), "value2");
+        const auto key1 = Sink::Storage::Key(Sink::Storage::Identifier::fromDisplayByteArray(uid), 6);
+        db.write(key1.toInternalByteArray(), "value1");
+        const auto key2 = Sink::Storage::Key(Sink::Storage::Identifier::fromDisplayByteArray(uid), 10);
+        db.write(key2.toInternalByteArray(), "value2");
         db.findLatest(uid, [&](const QByteArray &key, const QByteArray &value) { result = value; });
         QCOMPARE(result, QByteArray("value2"));
     }
@@ -689,7 +691,7 @@ private slots:
         Sink::Storage::DataStore::clearEnv();
 
         //Try to read-only dynamic opening of the db.
-        //This is the case if we don't have all databases available upon initializatoin and we don't (e.g. because the db hasn't been created yet) 
+        //This is the case if we don't have all databases available upon initializatoin and we don't (e.g. because the db hasn't been created yet)
         {
             // Trick the db into not loading all dbs by passing in a bogus layout.
             Sink::Storage::DataStore store(testDataPath, {dbName, {{"bogus", 0}}}, Sink::Storage::DataStore::ReadOnly);
@@ -707,7 +709,7 @@ private slots:
 
         Sink::Storage::DataStore::clearEnv();
         //Try to read-write dynamic opening of the db.
-        //This is the case if we don't have all databases available upon initializatoin and we don't (e.g. because the db hasn't been created yet) 
+        //This is the case if we don't have all databases available upon initializatoin and we don't (e.g. because the db hasn't been created yet)
         {
             // Trick the db into not loading all dbs by passing in a bogus layout.
             Sink::Storage::DataStore store(testDataPath, {dbName, {{"bogus", 0}}}, Sink::Storage::DataStore::ReadWrite);