(db/snapshot) Add snapshot DB unittest

core/src/db/meta/MetaAdapter.h

@@ -37,9 +37,9 @@ class MetaAdapter {
  public:
     MetaAdapter() {
         engine_ = std::make_shared<MockMetaEngine>();
-        //        DBMetaOptions options;
-        //        options.backend_uri_ = "mysql://root:12345678@127.0.0.1:3307/milvus";
-        //        engine_ = std::make_shared<MySqlEngine>(options);
+        //                DBMetaOptions options;
+        //                options.backend_uri_ = "mysql://root:12345678@127.0.0.1:3307/milvus";
+        //                engine_ = std::make_shared<MySqlEngine>(options);
     }
 
     SessionPtr
@@ -89,6 +89,8 @@ class MetaAdapter {
     template <typename ResourceT>
     Status
     Apply(snapshot::ResourceContextPtr<ResourceT> resp, int64_t& result_id) {
+        result_id = 0;
+
         auto session = CreateSession();
         session->Apply<ResourceT>(resp);
 

core/src/db/meta/backend/MockMetaEngine.cpp

@@ -73,7 +73,7 @@ MockMetaEngine::QueryNoLock(const MetaQueryContext& context, AttrsMapList& attrs
 }
 
 Status
-MockMetaEngine::AddNoLock(const MetaApplyContext& add_context, int64_t& result_id) {
+MockMetaEngine::AddNoLock(const MetaApplyContext& add_context, int64_t& result_id, TableRaw& pre_raw) {
     if (max_ip_map_.find(add_context.table_) == max_ip_map_.end() ||
         resources_.find(add_context.table_) == resources_.end()) {
         max_ip_map_[add_context.table_] = 0;
@@ -90,22 +90,24 @@ MockMetaEngine::AddNoLock(const MetaApplyContext& add_context, int64_t& result_i
 
     new_raw[F_ID] = std::to_string(max_id + 1);
     resources_[add_context.table_].push_back(new_raw);
+    pre_raw = new_raw;
     result_id = max_id + 1;
 
     return Status::OK();
 }
 
 Status
-MockMetaEngine::UpdateNoLock(const MetaApplyContext& update_context, int64_t& retult_id) {
+MockMetaEngine::UpdateNoLock(const MetaApplyContext& update_context, int64_t& result_id, TableRaw& pre_raw) {
     const std::string id_str = std::to_string(update_context.id_);
 
     auto& target_collection = resources_[update_context.table_];
     for (auto& attrs : target_collection) {
         if (attrs[F_ID] == id_str) {
+            pre_raw = attrs;
             for (auto& kv : update_context.attrs_) {
                 attrs[kv.first] = kv.second;
             }
-            retult_id = update_context.id_;
+            result_id = update_context.id_;
             return Status::OK();
         }
     }
@@ -115,12 +117,14 @@ MockMetaEngine::UpdateNoLock(const MetaApplyContext& update_context, int64_t& re
 }
 
 Status
-MockMetaEngine::DeleteNoLock(const MetaApplyContext& delete_context, int64_t& retult_id) {
+MockMetaEngine::DeleteNoLock(const MetaApplyContext& delete_context, int64_t& result_id, TableRaw& pre_raw) {
     const std::string id_str = std::to_string(delete_context.id_);
     auto& target_collection = resources_[delete_context.table_];
 
     for (auto iter = target_collection.begin(); iter != target_collection.end(); iter++) {
         if ((*iter)[F_ID] == id_str) {
+            pre_raw = *iter;
+            result_id = std::stol(iter->at(F_ID));
             target_collection.erase(iter);
             return Status::OK();
         }
@@ -141,18 +145,17 @@ MockMetaEngine::ExecuteTransaction(const std::vector<MetaApplyContext>& sql_cont
                                    std::vector<int64_t>& result_ids) {
     std::unique_lock<std::mutex> lock(mutex_);
 
-    auto duplicated_id_map = max_ip_map_;
-    auto duplicated_resource = resources_;
-
     auto status = Status::OK();
+    std::vector<std::pair<MetaContextOp, TableEntity>> pair_entities;
+    TableRaw raw;
     for (auto& context : sql_contexts) {
         int64_t id;
         if (context.op_ == oAdd) {
-            status = AddNoLock(context, id);
+            status = AddNoLock(context, id, raw);
         } else if (context.op_ == oUpdate) {
-            status = UpdateNoLock(context, id);
+            status = UpdateNoLock(context, id, raw);
         } else if (context.op_ == oDelete) {
-            status = DeleteNoLock(context, id);
+            status = DeleteNoLock(context, id, raw);
         } else {
             status = Status(SERVER_UNEXPECTED_ERROR, "Unknown resource context");
         }
@@ -161,14 +164,50 @@ MockMetaEngine::ExecuteTransaction(const std::vector<MetaApplyContext>& sql_cont
             break;
         }
         result_ids.push_back(id);
+        pair_entities.emplace_back(context.op_, TableEntity(context.table_, raw));
     }
 
     if (!status.ok()) {
-        max_ip_map_ = duplicated_id_map;
-        resources_ = duplicated_resource;
+        RollBackNoLock(pair_entities);
         return status;
     }
 
+    return status;
+}
+
+Status
+MockMetaEngine::RollBackNoLock(const std::vector<std::pair<MetaContextOp, TableEntity>>& pre_entities) {
+    for (auto& o_e : pre_entities) {
+        auto table = o_e.second.first;
+        if (o_e.first == oAdd) {
+            auto id = std::stol(o_e.second.second.at(F_ID));
+            max_ip_map_[table] = id - 1;
+            auto& table_res = resources_[table];
+            for (size_t i = 0; i < table_res.size(); i++) {
+                auto store_id = std::stol(table_res[i].at(F_ID));
+                if (store_id == id) {
+                    table_res.erase(table_res.begin() + i, table_res.begin() + i + 1);
+                    break;
+                }
+            }
+        } else if (o_e.first == oUpdate) {
+            auto id = std::stol(o_e.second.second.at(F_ID));
+            auto& table_res = resources_[table];
+            for (size_t j = 0; j < table_res.size(); j++) {
+                auto store_id = std::stol(table_res[j].at(F_ID));
+                if (store_id == id) {
+                    table_res.erase(table_res.begin() + j, table_res.begin() + j + 1);
+                    table_res.push_back(o_e.second.second);
+                    break;
+                }
+            }
+        } else if (o_e.first == oDelete) {
+            resources_[o_e.second.first].push_back(o_e.second.second);
+        } else {
+            continue;
+        }
+    }
+
     return Status::OK();
 }
 

core/src/db/meta/backend/MockMetaEngine.h

@@ -14,6 +14,7 @@
 #include <mutex>
 #include <string>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include "db/meta/backend/MetaEngine.h"
@@ -24,6 +25,7 @@ namespace milvus::engine::meta {
 class MockMetaEngine : public MetaEngine {
  private:
     using TableRaw = std::unordered_map<std::string, std::string>;
+    using TableEntity = std::pair<std::string, TableRaw>;
 
  public:
     MockMetaEngine() {
@@ -49,13 +51,16 @@ class MockMetaEngine : public MetaEngine {
     QueryNoLock(const MetaQueryContext& context, AttrsMapList& attrs);
 
     Status
-    AddNoLock(const MetaApplyContext& add_context, int64_t& retult_id);
+    AddNoLock(const MetaApplyContext& add_context, int64_t& result_id, TableRaw& pre_raw);
 
     Status
-    UpdateNoLock(const MetaApplyContext& add_context, int64_t& retult_id);
+    UpdateNoLock(const MetaApplyContext& add_context, int64_t& result_id, TableRaw& pre_raw);
 
     Status
-    DeleteNoLock(const MetaApplyContext& add_context, int64_t& retult_id);
+    DeleteNoLock(const MetaApplyContext& add_context, int64_t& result_id, TableRaw& pre_raw);
+
+    Status
+    RollBackNoLock(const std::vector<std::pair<MetaContextOp, TableEntity>>& pre_raws);
 
  private:
     std::mutex mutex_;

core/src/db/snapshot/ResourceContext.h

@@ -95,7 +95,8 @@ class ResourceContext {
 template <typename T>
 class ResourceContextBuilder {
  public:
-    ResourceContextBuilder() = default;
+    ResourceContextBuilder() : table_(T::Name), op_(meta::oAdd) {
+    }
 
     ResourceContextBuilder<T>&
     SetResource(typename T::Ptr res) {

core/unittest/ssdb/CMakeLists.txt

@@ -15,7 +15,8 @@ set(test_files
         ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/test_snapshot.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/test_segment.cpp
-        ${CMAKE_CURRENT_SOURCE_DIR}/test_db.cpp)
+        ${CMAKE_CURRENT_SOURCE_DIR}/test_db.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/test_ss_meta.cpp)
 
 add_executable(test_ssdb
         ${common_files}

core/unittest/ssdb/test_ss_meta.cpp

@@ -0,0 +1,172 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License.
+
+#include "ssdb/utils.h"
+
+#include "db/meta/MetaFields.h"
+#include "db/meta/backend/MetaContext.h"
+#include "db/snapshot/ResourceContext.h"
+
+template <typename T>
+using ResourceContext = milvus::engine::snapshot::ResourceContext<T>;
+template <typename T>
+using ResourceContextBuilder = milvus::engine::snapshot::ResourceContextBuilder<T>;
+
+using FType = milvus::engine::FieldType;
+using Op = milvus::engine::meta::MetaContextOp;
+
+TEST_F(SSMetaTest, ApplyTest) {
+    ID_TYPE result_id;
+
+    auto collection = std::make_shared<Collection>("meta_test_c1");
+    auto c_ctx = ResourceContextBuilder<Collection>().SetResource(collection).CreatePtr();
+    auto status = meta_.Apply<Collection>(c_ctx, result_id);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_GT(result_id, 0);
+    collection->SetID(result_id);
+
+    collection->Activate();
+    auto c2_ctx = ResourceContextBuilder<Collection>().SetResource(collection)
+        .SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
+    status = meta_.Apply<Collection>(c2_ctx, result_id);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_GT(result_id, 0);
+    ASSERT_EQ(result_id, collection->GetID());
+
+    auto c3_ctx = ResourceContextBuilder<Collection>().SetID(result_id).SetOp(Op::oDelete).CreatePtr();
+    status = meta_.Apply<Collection>(c3_ctx, result_id);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_GT(result_id, 0);
+    ASSERT_EQ(result_id, collection->GetID());
+}
+
+TEST_F(SSMetaTest, SessionTest) {
+    ID_TYPE result_id;
+
+    auto collection = std::make_shared<Collection>("meta_test_c1");
+    auto c_ctx = ResourceContextBuilder<Collection>().SetResource(collection).CreatePtr();
+    auto status = meta_.Apply<Collection>(c_ctx, result_id);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_GT(result_id, 0);
+    collection->SetID(result_id);
+
+    auto partition = std::make_shared<Partition>("meta_test_p1", result_id);
+    auto p_ctx = ResourceContextBuilder<Partition>().SetResource(partition).CreatePtr();
+    status = meta_.Apply<Partition>(p_ctx, result_id);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_GT(result_id, 0);
+    partition->SetID(result_id);
+
+    auto field = std::make_shared<Field>("meta_test_f1", 1, FType::INT64);
+    auto f_ctx = ResourceContextBuilder<Field>().SetResource(field).CreatePtr();
+    status = meta_.Apply<Field>(f_ctx, result_id);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_GT(result_id, 0);
+    field->SetID(result_id);
+
+    auto session = MetaAdapter::GetInstance().CreateSession();
+    ASSERT_TRUE(collection->Activate());
+    auto c2_ctx = ResourceContextBuilder<Collection>().SetResource(collection)
+        .SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
+    status = session->Apply<Collection>(c2_ctx);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+
+    ASSERT_TRUE(partition->Activate());
+    auto p2_ctx = ResourceContextBuilder<Partition>().SetResource(partition)
+        .SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
+    status = session->Apply<Partition>(p2_ctx);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+
+    ASSERT_TRUE(field->Activate());
+    auto f2_ctx = ResourceContextBuilder<Field>().SetResource(field)
+        .SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
+    status = session->Apply<Field>(f2_ctx);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+
+    std::vector<ID_TYPE> result_ids;
+    status = session->Commit(result_ids);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_EQ(result_ids.size(), 3);
+    ASSERT_EQ(result_ids.at(0), collection->GetID());
+    ASSERT_EQ(result_ids.at(1), partition->GetID());
+    ASSERT_EQ(result_ids.at(2), field->GetID());
+}
+
+TEST_F(SSMetaTest, SelectTest) {
+    ID_TYPE result_id;
+
+    auto collection = std::make_shared<Collection>("meta_test_c1");
+    ASSERT_TRUE(collection->Activate());
+    auto c_ctx = ResourceContextBuilder<Collection>().SetResource(collection).CreatePtr();
+    auto status = meta_.Apply<Collection>(c_ctx, result_id);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_GT(result_id, 0);
+    collection->SetID(result_id);
+
+    Collection::Ptr return_collection;
+    status = meta_.Select<Collection>(collection->GetID(), return_collection);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_EQ(collection->GetID(), return_collection->GetID());
+    ASSERT_EQ(collection->GetName(), return_collection->GetName());
+
+    auto collection2 = std::make_shared<Collection>("meta_test_c2");
+    ASSERT_TRUE(collection2->Activate());
+    auto c2_ctx = ResourceContextBuilder<Collection>().SetResource(collection2).CreatePtr();
+    status = meta_.Apply<Collection>(c2_ctx, result_id);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_GT(result_id, 0);
+    collection2->SetID(result_id);
+
+    ASSERT_GT(collection2->GetID(), collection->GetID());
+
+    std::vector<Collection::Ptr> return_collections;
+    status = meta_.SelectBy<Collection, ID_TYPE>(milvus::engine::meta::F_ID, collection2->GetID(), return_collections);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_EQ(return_collections.size(), 1);
+    ASSERT_EQ(return_collections.at(0)->GetID(), collection2->GetID());
+    ASSERT_EQ(return_collections.at(0)->GetName(), collection2->GetName());
+    return_collections.clear();
+
+    status = meta_.SelectBy<Collection, State>(milvus::engine::meta::F_STATE, State::ACTIVE, return_collections);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_EQ(return_collections.size(), 2);
+
+    std::vector<ID_TYPE> ids;
+    status = meta_.SelectResourceIDs<Collection, std::string>(ids, "", "");
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_EQ(ids.size(), 2);
+
+    ids.clear();
+    status = meta_.SelectResourceIDs<Collection, std::string>(ids, milvus::engine::meta::F_NAME, collection->GetName());
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_EQ(ids.size(), 1);
+    ASSERT_EQ(ids.at(0), collection->GetID());
+}
+
+TEST_F(SSMetaTest, TruncateTest) {
+    ID_TYPE result_id;
+
+    auto collection = std::make_shared<Collection>("meta_test_c1");
+    ASSERT_TRUE(collection->Activate());
+    auto c_ctx = ResourceContextBuilder<Collection>().SetResource(collection).CreatePtr();
+    auto status = meta_.Apply<Collection>(c_ctx, result_id);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_GT(result_id, 0);
+    collection->SetID(result_id);
+
+    status = meta_.TruncateAll();
+    ASSERT_TRUE(status.ok()) << status.ToString();
+
+    Collection::Ptr return_collection;
+    status = meta_.Select<Collection>(collection->GetID(), return_collection);
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    ASSERT_EQ(return_collection, nullptr);
+}