From a2735c8953b531e1d0f9bbf7b348862444225bc2 Mon Sep 17 00:00:00 2001
From: jbajic <jure.bajic@memgraph.com>
Date: Wed, 9 Nov 2022 12:10:46 +0100
Subject: [PATCH] Add label add/remove in UpdateVertex req

---
 src/query/v2/requests.hpp      |  8 +++---
 src/storage/v3/shard_rsm.cpp   | 47 +++++++++++++++++++++++++++++++++-
 tests/simulation/shard_rsm.cpp |  6 ++---
 3 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/src/query/v2/requests.hpp b/src/query/v2/requests.hpp
index 2ea4efb57..fe2851893 100644
--- a/src/query/v2/requests.hpp
+++ b/src/query/v2/requests.hpp
@@ -450,10 +450,12 @@ struct ExpandOneResponse {
   std::vector<ExpandOneResultRow> result;
 };
 
-struct UpdateVertexProp {
+struct UpdateVertex {
   PrimaryKey primary_key;
   // This should be a map
-  std::vector<std::pair<PropertyId, Value>> property_updates;
+  std::vector<LabelId> add_labels;
+  std::vector<LabelId> remove_labels;
+  std::map<PropertyId, Value> property_updates;
 };
 
 struct UpdateEdgeProp {
@@ -496,7 +498,7 @@ struct DeleteVerticesResponse {
 
 struct UpdateVerticesRequest {
   Hlc transaction_id;
-  std::vector<UpdateVertexProp> new_properties;
+  std::vector<UpdateVertex> update_vertices;
 };
 
 struct UpdateVerticesResponse {
diff --git a/src/storage/v3/shard_rsm.cpp b/src/storage/v3/shard_rsm.cpp
index f288e5e27..bf66bc463 100644
--- a/src/storage/v3/shard_rsm.cpp
+++ b/src/storage/v3/shard_rsm.cpp
@@ -535,12 +535,46 @@ msgs::WriteResponses ShardRsm::ApplyWrite(msgs::CreateVerticesRequest &&req) {
   return msgs::CreateVerticesResponse{.success = action_successful};
 }
 
+void HandleError(const ResultErrorType &error, const std::string_view action) {
+  std::visit(
+      [action]<typename T>(T &&error) {
+        using ErrorType = std::remove_cvref_t<T>;
+        if constexpr (std::is_same_v<ErrorType, SchemaViolation>) {
+          spdlog::debug("{} failed with error: SchemaViolation", action);
+        } else if constexpr (std::is_same_v<ErrorType, Error>) {
+          switch (error) {
+            case Error::DELETED_OBJECT:
+              spdlog::debug("{} failed with error: DELETED_OBJECT", action);
+              break;
+            case Error::NONEXISTENT_OBJECT:
+              spdlog::debug("{} failed with error: NONEXISTENT_OBJECT", action);
+              break;
+            case Error::SERIALIZATION_ERROR:
+              spdlog::debug("{} failed with error: SERIALIZATION_ERROR", action);
+              break;
+            case Error::PROPERTIES_DISABLED:
+              spdlog::debug("{} failed with error: PROPERTIES_DISABLED", action);
+              break;
+            case Error::VERTEX_HAS_EDGES:
+              spdlog::debug("{} failed with error: VERTEX_HAS_EDGES", action);
+              break;
+            case Error::VERTEX_ALREADY_INSERTED:
+              spdlog::debug("{} failed with error: VERTEX_ALREADY_INSERTED", action);
+              break;
+          }
+        } else {
+          static_assert(kAlwaysFalse<T>, "Missing type from variant visitor");
+        }
+      },
+      error);
+}
+
 msgs::WriteResponses ShardRsm::ApplyWrite(msgs::UpdateVerticesRequest &&req) {
   auto acc = shard_->Access(req.transaction_id);
 
   bool action_successful = true;
 
-  for (auto &vertex : req.new_properties) {
+  for (auto &vertex : req.update_vertices) {
     if (!action_successful) {
       break;
     }
@@ -553,6 +587,17 @@ msgs::WriteResponses ShardRsm::ApplyWrite(msgs::UpdateVerticesRequest &&req) {
       continue;
     }
 
+    for (const auto label : vertex.add_labels) {
+      if (const auto maybe_error = vertex_to_update->AddLabelAndValidate(label); maybe_error.HasError()) {
+        HandleError(maybe_error.GetError(), "Update Vertex");
+      }
+    }
+    for (const auto label : vertex.remove_labels) {
+      if (const auto maybe_error = vertex_to_update->RemoveLabelAndValidate(label); maybe_error.HasError()) {
+        HandleError(maybe_error.GetError(), "Update Vertex");
+      }
+    }
+
     for (auto &update_prop : vertex.property_updates) {
       auto result_schema =
           vertex_to_update->SetPropertyAndValidate(update_prop.first, ToPropertyValue(std::move(update_prop.second)));
diff --git a/tests/simulation/shard_rsm.cpp b/tests/simulation/shard_rsm.cpp
index 64d0a0861..23dfecac1 100644
--- a/tests/simulation/shard_rsm.cpp
+++ b/tests/simulation/shard_rsm.cpp
@@ -189,13 +189,13 @@ bool AttemptToUpdateVertex(ShardClient &client, int64_t value) {
   std::vector<std::pair<PropertyId, msgs::Value>> property_updates;
   auto property_update = std::make_pair(PropertyId::FromUint(5), msgs::Value(static_cast<int64_t>(10000)));
 
-  auto vertex_prop = msgs::UpdateVertexProp{};
+  msgs::UpdateVertex vertex_prop;
   vertex_prop.primary_key = vertex_id;
   vertex_prop.property_updates = {property_update};
 
-  auto update_req = msgs::UpdateVerticesRequest{};
+  msgs::UpdateVerticesRequest update_req;
   update_req.transaction_id.logical_id = GetTransactionId();
-  update_req.new_properties = {vertex_prop};
+  update_req.update_vertices = {vertex_prop};
 
   while (true) {
     auto write_res = client.SendWriteRequest(update_req);