From da288af368625f96a34167edff8a02bbfdd25c86 Mon Sep 17 00:00:00 2001
From: Muhammad Hamza Sajjad
 <45758804+muhammadhamzasajjad@users.noreply.github.com>
Date: Thu, 8 Feb 2024 12:40:55 +0000
Subject: [PATCH] #447 Memory Storage Exception Normalization (#1297)

Implementation of memory storage specific exceptions normalization. This
PR is a refactor of #584.
---
 cpp/arcticdb/storage/memory/memory_storage.cpp    |  7 +++++--
 cpp/arcticdb/storage/storage.hpp                  |  7 +++++++
 .../storage/test/test_storage_exceptions.cpp      | 15 ++++++++++++++-
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/cpp/arcticdb/storage/memory/memory_storage.cpp b/cpp/arcticdb/storage/memory/memory_storage.cpp
index 8a2a9bd953..b80fed3ba9 100644
--- a/cpp/arcticdb/storage/memory/memory_storage.cpp
+++ b/cpp/arcticdb/storage/memory/memory_storage.cpp
@@ -58,7 +58,10 @@ namespace arcticdb::storage::memory {
             for (auto &kv : group.values()) {
                 auto it = key_vec.find(kv.variant_key());
 
-                util::check_rte(opts.upsert_ || it != key_vec.end(), "update called with upsert=false but key does not exist");
+                if (!opts.upsert_ && it == key_vec.end()) {
+                    std::string err_message = fmt::format("do_update called with upsert=false on non-existent key(s): {}", kv.variant_key());
+                    throw KeyNotFoundException(std::move(kv.variant_key()), err_message);
+                }
 
                 if(it != key_vec.end()) {
                     key_vec.erase(it);
@@ -110,7 +113,7 @@ namespace arcticdb::storage::memory {
                     ARCTICDB_DEBUG(log::storage(), "Read key {}: {}, with {} bytes of data", variant_key_type(k), variant_key_view(k));
                     key_vec.erase(it);
                 } else if (!opts.ignores_missing_key_) {
-                    util::raise_rte("Failed to find segment for key {}",variant_key_view(k));
+                    throw KeyNotFoundException(std::move(k));
                 }
             }
         });
diff --git a/cpp/arcticdb/storage/storage.hpp b/cpp/arcticdb/storage/storage.hpp
index d51d723944..17ffb0378f 100644
--- a/cpp/arcticdb/storage/storage.hpp
+++ b/cpp/arcticdb/storage/storage.hpp
@@ -73,6 +73,13 @@ class KeyNotFoundException : public ArcticSpecificException<ErrorCode::E_KEY_NOT
             keys_(std::make_shared<Composite<VariantKey>>(std::move(keys))) {
     }
 
+    explicit KeyNotFoundException(const VariantKey& single_key):
+            KeyNotFoundException(Composite<VariantKey>{VariantKey{single_key}}) {}
+
+    explicit KeyNotFoundException(const VariantKey& single_key, std::string err_output):
+            KeyNotFoundException(Composite<VariantKey>{VariantKey{single_key}}, err_output) {}
+
+
     Composite<VariantKey>& keys() {
         return *keys_;
     }
diff --git a/cpp/arcticdb/storage/test/test_storage_exceptions.cpp b/cpp/arcticdb/storage/test/test_storage_exceptions.cpp
index 9211e4ade4..77a258e5fa 100644
--- a/cpp/arcticdb/storage/test/test_storage_exceptions.cpp
+++ b/cpp/arcticdb/storage/test/test_storage_exceptions.cpp
@@ -10,6 +10,7 @@
 #include <arcticdb/codec/codec.hpp>
 #include <arcticdb/storage/storage.hpp>
 #include <arcticdb/storage/lmdb/lmdb_storage.hpp>
+#include <arcticdb/storage/memory/memory_storage.hpp>
 #include <arcticdb/util/buffer.hpp>
 
 #include <filesystem>
@@ -61,6 +62,15 @@ class LMDBStorageFactory : public StorageFactory {
     }
 };
 
+class MemoryStorageFactory : public StorageFactory {
+    std::unique_ptr<arcticdb::storage::Storage> create() override {
+        arcticdb::proto::memory_storage::Config cfg;
+        arcticdb::storage::LibraryPath library_path{"a", "b"};
+
+        return std::make_unique<arcticdb::storage::memory::MemoryStorage>(library_path, arcticdb::storage::OpenMode::WRITE, cfg);
+    }
+};
+
 // Generic tests that run with all types of storages
 
 class GenericStorageTest : public ::testing::TestWithParam<std::shared_ptr<StorageFactory>> {
@@ -134,7 +144,10 @@ TEST_P(GenericStorageTest, RemoveKeyNotFoundException) {
 INSTANTIATE_TEST_SUITE_P(
         AllStoragesCommonTests,
         GenericStorageTest,
-        ::testing::Values(std::make_shared<LMDBStorageFactory>())
+        ::testing::Values(
+                std::make_shared<LMDBStorageFactory>(),
+                std::make_shared<MemoryStorageFactory>()
+        )
 );
 
 // LMDB Storage specific tests