Fix unit test

src/dbms/dbms_handler.cpp

@@ -10,6 +10,7 @@
 // licenses/APL.txt.
 
 #include "dbms/dbms_handler.hpp"
+#include "utils/exceptions.hpp"
 
 namespace memgraph::dbms {
 #ifdef MG_ENTERPRISE
@@ -52,8 +53,8 @@ DbmsHandler::DbmsHandler(
   // Startup replication state (if recovered at startup)
   auto replica = [this](replication::RoleReplicaData const &data) {
     // Register handlers
-    InMemoryReplicationHandlers::Register(this, *data.server);
-    if (!data.server->Start()) {
+    InMemoryReplicationHandlers::Register(this, *data.server_);
+    if (!data.server_->Start()) {
       spdlog::error("Unable to start the replication server.");
       return false;
     }
@@ -66,8 +67,11 @@ DbmsHandler::DbmsHandler(
     }
     return true;
   };
+  auto coordinator = [](replication::RoleCoordinatorData &) -> bool {
+    throw utils::NotYetImplemented("Not yet implemented");
+  };
   // Startup proccess for main/replica
-  MG_ASSERT(std::visit(memgraph::utils::Overloaded{replica, main}, repl_state_.ReplicationData()),
+  MG_ASSERT(std::visit(memgraph::utils::Overloaded{replica, main, coordinator}, repl_state_.ReplicationData()),
             "Replica recovery failure!");
 }
 #endif

src/dbms/replication_handler.cpp

@@ -18,7 +18,7 @@
 #include "dbms/replication_client.hpp"
 #include "replication/state.hpp"
 
-using memgraph::replication::ReplicationClientConfig;
+using memgraph::replication::ReplicationServer;
 using memgraph::replication::ReplicationState;
 using memgraph::replication::RoleMainData;
 using memgraph::replication::RoleReplicaData;
@@ -44,12 +44,33 @@ std::string RegisterReplicaErrorToString(RegisterReplicaError error) {
 
 ReplicationHandler::ReplicationHandler(DbmsHandler &dbms_handler) : dbms_handler_(dbms_handler) {}
 
+#ifdef MG_ENTERPRISE
+bool ReplicationHandler::SetReplicationRoleMain(const memgraph::replication::ReplicationServerConfig &config) {
+#else
 bool ReplicationHandler::SetReplicationRoleMain() {
+#endif
+
+#ifdef MG_ENTERPRISE
+  auto const main_handler = [&config](RoleMainData &main_data) {
+    // TODO: (andi) Epoch will probably need to updated eventually, maybe with coordinator command.
+    // If we are already MAIN, we will not update epoch. Instead, we wil set the config and create ReplicationServer.
+    // Registered replicas won't be touched.
+    main_data.server_config_ = config;
+    main_data.server_ = std::make_unique<ReplicationServer>(config);
+    return true;
+  };
+#else
   auto const main_handler = [](RoleMainData const &) {
     // If we are already MAIN, we don't want to change anything
     return false;
   };
+#endif
+
+#ifdef MG_ENTERPRISE
+  auto const replica_handler = [this, &config](RoleReplicaData const &) {
+#else
   auto const replica_handler = [this](RoleReplicaData const &) {
+#endif
     // STEP 1) bring down all REPLICA servers
     dbms_handler_.ForEach([](Database *db) {
       auto *storage = db->storage();
@@ -57,9 +78,13 @@ bool ReplicationHandler::SetReplicationRoleMain() {
       storage->PrepareForNewEpoch();
     });
 
-    // STEP 2) Change to MAIN
-    // TODO: restore replication servers if false?
+// STEP 2) Change to MAIN
+// TODO: restore replication servers if false?
+#ifdef MG_ENTERPRISE
+    if (!dbms_handler_.ReplicationState().SetReplicationRoleMain(config)) {
+#else
     if (!dbms_handler_.ReplicationState().SetReplicationRoleMain()) {
+#endif
       // TODO: Handle recovery on failure???
       return false;
     }
@@ -75,17 +100,50 @@ bool ReplicationHandler::SetReplicationRoleMain() {
     return true;
   };
 
+#ifdef MG_ENTERPRISE
+  // COORDINATOR cannot become main
+  auto const coordinator_handler = [](replication::RoleCoordinatorData const &) { return false; };
+  // TODO: under lock
+  return std::visit(utils::Overloaded{main_handler, replica_handler, coordinator_handler},
+                    dbms_handler_.ReplicationState().ReplicationData());
+#else
   // TODO: under lock
   return std::visit(utils::Overloaded{main_handler, replica_handler},
                     dbms_handler_.ReplicationState().ReplicationData());
+#endif
 }
 
+#ifdef MG_ENTERPRISE
+bool ReplicationHandler::SetReplicationRoleCoordinator() {
+  // Upgrading REPLICA to COORDINATOR is not supported
+  auto const replica_handler = [](RoleReplicaData const &) { return false; };
+  // If we are already COORDINATOR, we don't want to change anything
+  auto const coordinator_handler = [](replication::RoleCoordinatorData const &) { return false; };
+  // Upgrade MAIN to COORDINATOR
+  // TODO: (andi) Probably more complex steps will be necessary here.
+  auto const main_handler = [this](RoleMainData const &) {
+    return dbms_handler_.ReplicationState().SetReplicationRoleCoordinator();
+  };
+
+  // TODO: under lock
+  return std::visit(utils::Overloaded{main_handler, replica_handler, coordinator_handler},
+                    dbms_handler_.ReplicationState().ReplicationData());
+}
+#endif
+
 bool ReplicationHandler::SetReplicationRoleReplica(const memgraph::replication::ReplicationServerConfig &config) {
   // We don't want to restart the server if we're already a REPLICA
   if (dbms_handler_.ReplicationState().IsReplica()) {
     return false;
   }
 
+#ifdef MG_ENTERPRISE
+  // Upgrading replica to coordinator is forbidden
+  if (dbms_handler_.ReplicationState().IsCoordinator()) {
+    return false;
+  }
+#endif
+
   // TODO StorageState needs to be synched. Could have a dangling reference if someone adds a database as we are
   //      deleting the replica.
   // Remove database specific clients
@@ -99,34 +157,40 @@ bool ReplicationHandler::SetReplicationRoleReplica(const memgraph::replication::
   // Creates the server
   dbms_handler_.ReplicationState().SetReplicationRoleReplica(config);
 
-  // Start
+  // TODO: (andi) Assert instead of returning false
+  auto const main_handler = [](RoleMainData const &) { return false; };
+
+  auto const replica_handler = [this](RoleReplicaData const &data) {
+    InMemoryReplicationHandlers::Register(&dbms_handler_, *data.server_);
+    if (!data.server_->Start()) {
+      spdlog::error("Unable to start the replication server.");
+      return false;
+    }
+    return true;
+  };
+
+#ifdef MG_ENTERPRISE
+  // TODO: (andi) Assert instead of returning false
+  auto const coordinator_handler = [](replication::RoleCoordinatorData const &) { return false; };
+  const auto success = std::visit(utils::Overloaded{main_handler, replica_handler, coordinator_handler},
+                                  dbms_handler_.ReplicationState().ReplicationData());
+#else
   const auto success =
-      std::visit(utils::Overloaded{[](RoleMainData const &) {
-                                     // ASSERT
-                                     return false;
-                                   },
-                                   [this](RoleReplicaData const &data) {
-                                     // Register handlers
-                                     InMemoryReplicationHandlers::Register(&dbms_handler_, *data.server);
-                                     if (!data.server->Start()) {
-                                       spdlog::error("Unable to start the replication server.");
-                                       return false;
-                                     }
-                                     return true;
-                                   }},
-                 dbms_handler_.ReplicationState().ReplicationData());
+      std::visit(utils::Overloaded{main_handler, replica_handler}, dbms_handler_.ReplicationState().ReplicationData());
+#endif
+
   // TODO Handle error (restore to main?)
   return success;
 }
 
 auto ReplicationHandler::RegisterReplica(const memgraph::replication::ReplicationClientConfig &config)
     -> memgraph::utils::BasicResult<RegisterReplicaError> {
-  MG_ASSERT(dbms_handler_.ReplicationState().IsMain(), "Only main instance can register a replica!");
+  MG_ASSERT(!dbms_handler_.ReplicationState().IsReplica(), "Replica can't register another replica!");
 
   auto instance_client = dbms_handler_.ReplicationState().RegisterReplica(config);
   if (instance_client.HasError()) switch (instance_client.GetError()) {
-      case memgraph::replication::RegisterReplicaError::NOT_MAIN:
-        MG_ASSERT(false, "Only main instance can register a replica!");
+      case memgraph::replication::RegisterReplicaError::IS_REPLICA:
+        MG_ASSERT(false, "Replica can't register another replica!");
         return {};
       case memgraph::replication::RegisterReplicaError::NAME_EXISTS:
         return memgraph::dbms::RegisterReplicaError::NAME_EXISTS;
@@ -180,28 +244,83 @@ auto ReplicationHandler::RegisterReplica(const memgraph::replication::Replicatio
   return {};
 }
 
+#ifdef MG_ENTERPRISE
+auto ReplicationHandler::RegisterMain(const memgraph::replication::ReplicationClientConfig &config)
+    -> utils::BasicResult<RegisterMainError> {
+  MG_ASSERT(dbms_handler_.ReplicationState().IsCoordinator(), "Only coordinator can register main instance!");
+
+  auto instance_client = dbms_handler_.ReplicationState().RegisterMain(config);
+  if (instance_client.HasError()) switch (instance_client.GetError()) {
+      case memgraph::replication::RegisterMainError::NOT_COORDINATOR:
+        MG_ASSERT(false, "Only coordinator can register main instance!");
+        return {};
+      case memgraph::replication::RegisterMainError::END_POINT_EXISTS:
+        return memgraph::dbms::RegisterMainError::END_POINT_EXISTS;
+      case memgraph::replication::RegisterMainError::COULD_NOT_BE_PERSISTED:
+        return memgraph::dbms::RegisterMainError::COULD_NOT_BE_PERSISTED;
+      case memgraph::replication::RegisterMainError::MAIN_ALREADY_EXISTS:
+        return memgraph::dbms::RegisterMainError::MAIN_ALREADY_EXISTS;
+      case memgraph::replication::RegisterMainError::SUCCESS:
+        break;
+    }
+
+  if (!allow_mt_repl && dbms_handler_.All().size() > 1) {
+    spdlog::warn("Multi-tenant replication is currently not supported!");
+  }
+
+  return {};
+}
+#endif
+
 auto ReplicationHandler::UnregisterReplica(std::string_view name) -> UnregisterReplicaResult {
   auto const replica_handler = [](RoleReplicaData const &) -> UnregisterReplicaResult {
-    return UnregisterReplicaResult::NOT_MAIN;
+    return UnregisterReplicaResult::IS_REPLICA;
   };
+
   auto const main_handler = [this, name](RoleMainData &mainData) -> UnregisterReplicaResult {
-    if (!dbms_handler_.ReplicationState().TryPersistUnregisterReplica(name)) {
+    if (!dbms_handler_.ReplicationState().TryPersistUnregisterReplicaOnMain(name)) {
       return UnregisterReplicaResult::COULD_NOT_BE_PERSISTED;
     }
+
     // Remove database specific clients
     dbms_handler_.ForEach([name](Database *db) {
       db->storage()->repl_storage_state_.replication_clients_.WithLock([&name](auto &clients) {
         std::erase_if(clients, [name](const auto &client) { return client->Name() == name; });
       });
     });
+
     // Remove instance level clients
     auto const n_unregistered =
         std::erase_if(mainData.registered_replicas_, [name](auto const &client) { return client.name_ == name; });
     return n_unregistered != 0 ? UnregisterReplicaResult::SUCCESS : UnregisterReplicaResult::CAN_NOT_UNREGISTER;
   };
 
+#ifdef MG_ENTERPRISE
+  auto const coordinator_handler =
+      [this, name](replication::RoleCoordinatorData &coordinatorData) -> UnregisterReplicaResult {
+    if (!dbms_handler_.ReplicationState().TryPersistUnregisterReplicaOnMain(name)) {
+      return UnregisterReplicaResult::COULD_NOT_BE_PERSISTED;
+    }
+
+    // Remove database specific clients
+    dbms_handler_.ForEach([name](Database *db) {
+      db->storage()->repl_storage_state_.replication_clients_.WithLock([&name](auto &clients) {
+        std::erase_if(clients, [name](const auto &client) { return client->Name() == name; });
+      });
+    });
+
+    // Remove instance level clients
+    auto const n_unregistered = std::erase_if(coordinatorData.registered_replicas_,
+                                              [name](auto const &client) { return client.name_ == name; });
+    return n_unregistered != 0 ? UnregisterReplicaResult::SUCCESS : UnregisterReplicaResult::CAN_NOT_UNREGISTER;
+  };
+
+  return std::visit(utils::Overloaded{main_handler, replica_handler, coordinator_handler},
+                    dbms_handler_.ReplicationState().ReplicationData());
+#else
   return std::visit(utils::Overloaded{main_handler, replica_handler},
                     dbms_handler_.ReplicationState().ReplicationData());
+#endif
 }
 
 auto ReplicationHandler::GetRole() const -> memgraph::replication::ReplicationRole {
@@ -212,12 +331,15 @@ bool ReplicationHandler::IsMain() const { return dbms_handler_.ReplicationState(
 
 bool ReplicationHandler::IsReplica() const { return dbms_handler_.ReplicationState().IsReplica(); }
 
+#ifdef MG_ENTERPRISE
+bool ReplicationHandler::IsCoordinator() const { return dbms_handler_.ReplicationState().IsCoordinator(); }
+#endif
+
 // Per storage
 // NOTE Storage will connect to all replicas. Future work might change this
 void RestoreReplication(replication::ReplicationState &repl_state, storage::Storage &storage) {
   spdlog::info("Restoring replication role.");
 
-  /// MAIN
   auto const recover_main = [&storage](RoleMainData &mainData) {
     // Each individual client has already been restored and started. Here we just go through each database and start its
     // client
@@ -249,14 +371,15 @@ void RestoreReplication(replication::ReplicationState &repl_state, storage::Stor
     spdlog::info("Replication role restored to MAIN.");
   };
 
-  /// REPLICA
   auto const recover_replica = [](RoleReplicaData const &data) { /*nothing to do*/ };
 
-  std::visit(
-      utils::Overloaded{
-          recover_main,
-          recover_replica,
-      },
-      repl_state.ReplicationData());
+#ifdef MG_ENTERPRISE
+  // TODO: (andi) This will probably have to be implemented
+  auto const recover_coordinator = [](replication::RoleCoordinatorData const &) { /*nothing to do*/ };
+  std::visit(utils::Overloaded{recover_main, recover_replica, recover_coordinator}, repl_state.ReplicationData());
+#else
+  std::visit(utils::Overloaded{recover_main, recover_replica}, repl_state.ReplicationData());
+#endif
 }
+
 }  // namespace memgraph::dbms

src/dbms/replication_handler.hpp

@@ -25,9 +25,15 @@ struct ReplicationClientConfig;
 namespace memgraph::dbms {
 class DbmsHandler;
 
+/// TODO: (andi) Two definitions of the same enum
 enum class RegisterReplicaError : uint8_t { NAME_EXISTS, END_POINT_EXISTS, CONNECTION_FAILED, COULD_NOT_BE_PERSISTED };
+
+#ifdef MG_ENTERPRISE
+enum class RegisterMainError : uint8_t { MAIN_ALREADY_EXISTS, END_POINT_EXISTS, COULD_NOT_BE_PERSISTED };
+#endif
+
 enum class UnregisterReplicaResult : uint8_t {
-  NOT_MAIN,
+  IS_REPLICA,
   COULD_NOT_BE_PERSISTED,
   CAN_NOT_UNREGISTER,
   SUCCESS,
@@ -38,16 +44,33 @@ enum class UnregisterReplicaResult : uint8_t {
 struct ReplicationHandler {
   explicit ReplicationHandler(DbmsHandler &dbms_handler);
 
-  // as REPLICA, become MAIN
+#ifdef MG_ENTERPRISE
+  // as default main, add replication server to the main.
+  // As replica become main
+  bool SetReplicationRoleMain(const memgraph::replication::ReplicationServerConfig &config);
+#else
+  // as replica, become main.
   bool SetReplicationRoleMain();
+#endif
 
-  // as MAIN, become REPLICA
+  // as main, become replica
   bool SetReplicationRoleReplica(const memgraph::replication::ReplicationServerConfig &config);
 
-  // as MAIN, define and connect to REPLICAs
+#ifdef MG_ENTERPRISE
+  // as default main, become coordinator
+  bool SetReplicationRoleCoordinator();
+#endif
+
+  // as main, define and connect to replicas
   auto RegisterReplica(const memgraph::replication::ReplicationClientConfig &config)
       -> utils::BasicResult<RegisterReplicaError>;
 
+#ifdef MG_ENTERPRISE
+  // as coordinator, connect to main
+  auto RegisterMain(const memgraph::replication::ReplicationClientConfig &config)
+      -> utils::BasicResult<RegisterMainError>;
+#endif
+
   // as MAIN, remove a REPLICA connection
   auto UnregisterReplica(std::string_view name) -> UnregisterReplicaResult;
 
@@ -55,6 +78,9 @@ struct ReplicationHandler {
   auto GetRole() const -> memgraph::replication::ReplicationRole;
   bool IsMain() const;
   bool IsReplica() const;
+#ifdef MG_ENTERPRISE
+  bool IsCoordinator() const;
+#endif
 
  private:
   DbmsHandler &dbms_handler_;

src/glue/SessionHL.cpp

@@ -193,12 +193,13 @@ std::pair<std::vector<std::string>, std::optional<int>> SessionHL::Interpret(
   for (const auto &[key, bolt_param] : params) {
     params_pv.emplace(key, ToPropertyValue(bolt_param));
   }
+
+#ifdef MG_ENTERPRISE
   const std::string *username{nullptr};
   if (user_) {
     username = &user_->username();
   }
 
-#ifdef MG_ENTERPRISE
   // TODO: Update once interpreter can handle non-database queries (db_acc will be nullopt)
   auto *db = interpreter_.current_db_.db_acc_->get();
   if (memgraph::license::global_license_checker.IsEnterpriseValidFast()) {

src/query/frontend/ast/ast.hpp

@@ -3025,9 +3025,20 @@ class ReplicationQuery : public memgraph::query::Query {
   static const utils::TypeInfo kType;
   const utils::TypeInfo &GetTypeInfo() const override { return kType; }
 
-  enum class Action { SET_REPLICATION_ROLE, SHOW_REPLICATION_ROLE, REGISTER_REPLICA, DROP_REPLICA, SHOW_REPLICAS };
+  enum class Action {
+    SET_REPLICATION_ROLE,
+    SHOW_REPLICATION_ROLE,
+    REGISTER_REPLICA,
+    DROP_REPLICA,
+    SHOW_REPLICAS,
+    REGISTER_MAIN
+  };
 
+#ifdef MG_ENTERPRISE
+  enum class ReplicationRole{MAIN, REPLICA, COORDINATOR};
+#else
   enum class ReplicationRole { MAIN, REPLICA };
+#endif
 
   enum class SyncMode { SYNC, ASYNC };
 

src/query/frontend/ast/cypher_main_visitor.cpp

@@ -316,6 +316,30 @@ antlrcpp::Any CypherMainVisitor::visitEdgeImportModeQuery(MemgraphCypher::EdgeIm
 antlrcpp::Any CypherMainVisitor::visitSetReplicationRole(MemgraphCypher::SetReplicationRoleContext *ctx) {
   auto *replication_query = storage_->Create<ReplicationQuery>();
   replication_query->action_ = ReplicationQuery::Action::SET_REPLICATION_ROLE;
+
+#ifdef MG_ENTERPRISE
+  // Licence check in the interpreter.
+  if (ctx->MAIN() || ctx->REPLICA()) {
+    if (!ctx->WITH() || !ctx->PORT()) {
+      throw SemanticException("Port must be specified when setting replication role to main or replica!");
+    }
+    if (ctx->port->numberLiteral() && ctx->port->numberLiteral()->integerLiteral()) {
+      replication_query->port_ = std::any_cast<Expression *>(ctx->port->accept(this));
+    } else {
+      throw SyntaxException("Port must be an integer literal!");
+    }
+  } else if (ctx->WITH() || ctx->PORT()) {  // coordinator
+    throw SemanticException("Port shouldn't be specified when setting replication role to coordinator!");
+  }
+
+  if (ctx->MAIN()) {
+    replication_query->role_ = ReplicationQuery::ReplicationRole::MAIN;
+  } else if (ctx->REPLICA()) {
+    replication_query->role_ = ReplicationQuery::ReplicationRole::REPLICA;
+  } else if (ctx->COORDINATOR()) {
+    replication_query->role_ = ReplicationQuery::ReplicationRole::COORDINATOR;
+  }
+#else
   if (ctx->MAIN()) {
     if (ctx->WITH() || ctx->PORT()) {
       throw SemanticException("Main can't set a port!");
@@ -331,6 +355,8 @@ antlrcpp::Any CypherMainVisitor::visitSetReplicationRole(MemgraphCypher::SetRepl
       }
     }
   }
+#endif
+
   return replication_query;
 }
 antlrcpp::Any CypherMainVisitor::visitShowReplicationRole(MemgraphCypher::ShowReplicationRoleContext *ctx) {
@@ -351,9 +377,21 @@ antlrcpp::Any CypherMainVisitor::visitRegisterReplica(MemgraphCypher::RegisterRe
 
   if (!ctx->socketAddress()->literal()->StringLiteral()) {
     throw SemanticException("Socket address should be a string literal!");
-  } else {
-    replication_query->socket_address_ = std::any_cast<Expression *>(ctx->socketAddress()->accept(this));
   }
+  replication_query->socket_address_ = std::any_cast<Expression *>(ctx->socketAddress()->accept(this));
+
+  return replication_query;
+}
+
+// We want to completely disable this code in community version.
+// License check is done in the interpreter.
+antlrcpp::Any CypherMainVisitor::visitRegisterMain(MemgraphCypher::RegisterMainContext *ctx) {
+  auto *replication_query = storage_->Create<ReplicationQuery>();
+  replication_query->action_ = ReplicationQuery::Action::REGISTER_MAIN;
+  if (!ctx->socketAddress()->literal()->StringLiteral()) {
+    throw SemanticException("Socket address should be a string literal!");
+  }
+  replication_query->socket_address_ = std::any_cast<Expression *>(ctx->socketAddress()->accept(this));
 
   return replication_query;
 }

src/query/frontend/ast/cypher_main_visitor.hpp

@@ -221,6 +221,11 @@ class CypherMainVisitor : public antlropencypher::MemgraphCypherBaseVisitor {
    */
   antlrcpp::Any visitRegisterReplica(MemgraphCypher::RegisterReplicaContext *ctx) override;
 
+  /**
+   * @return ReplicationQuery*
+   */
+  antlrcpp::Any visitRegisterMain(MemgraphCypher::RegisterMainContext *ctx) override;
+
   /**
    * @return ReplicationQuery*
    */

src/query/frontend/opencypher/grammar/MemgraphCypher.g4

@@ -42,6 +42,7 @@ memgraphCypherKeyword : cypherKeyword
                       | CONSUMER_GROUP
                       | CREATE_DELETE
                       | CREDENTIALS
+                      | COORDINATOR
                       | CSV
                       | DATA
                       | DELIMITER
@@ -179,6 +180,7 @@ authQuery : createRole
 replicationQuery : setReplicationRole
                  | showReplicationRole
                  | registerReplica
+                 | registerMain
                  | dropReplica
                  | showReplicas
                  ;
@@ -359,7 +361,7 @@ dumpQuery : DUMP DATABASE ;
 
 analyzeGraphQuery : ANALYZE GRAPH ( ON LABELS ( listOfColonSymbolicNames | ASTERISK ) ) ? ( DELETE STATISTICS ) ? ;
 
-setReplicationRole : SET REPLICATION ROLE TO ( MAIN | REPLICA )
+setReplicationRole : SET REPLICATION ROLE TO ( MAIN | REPLICA | COORDINATOR )
                       ( WITH PORT port=literal ) ? ;
 
 showReplicationRole : SHOW REPLICATION ROLE ;
@@ -371,6 +373,8 @@ socketAddress : literal ;
 registerReplica : REGISTER REPLICA replicaName ( SYNC | ASYNC )
                 TO socketAddress ;
 
+registerMain : REGISTER MAIN TO socketAddress ;
+
 dropReplica : DROP REPLICA replicaName ;
 
 showReplicas : SHOW REPLICAS ;

src/query/frontend/opencypher/grammar/MemgraphCypherLexer.g4

@@ -47,6 +47,8 @@ CONSUMER_GROUP          : C O N S U M E R UNDERSCORE G R O U P ;
 CREATE_DELETE           : C R E A T E UNDERSCORE D E L E T E ;
 CREDENTIALS             : C R E D E N T I A L S ;
 CSV                     : C S V ;
+COORDINATOR             : C O O R D I N A T O R;
+CLUSTER                 : C L U S T E R;
 DATA                    : D A T A ;
 DELIMITER               : D E L I M I T E R ;
 DATABASE                : D A T A B A S E ;

src/query/interpreter.cpp

@@ -278,22 +278,55 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
 
   /// @throw QueryRuntimeException if an error ocurred.
   void SetReplicationRole(ReplicationQuery::ReplicationRole replication_role, std::optional<int64_t> port) override {
-    if (replication_role == ReplicationQuery::ReplicationRole::MAIN) {
-      if (!handler_.SetReplicationRoleMain()) {
-        throw QueryRuntimeException("Couldn't set role to main!");
-      }
-    } else {
+    auto ValidatePort = [](std::optional<int64_t> port) -> void {
       if (!port || *port < 0 || *port > std::numeric_limits<uint16_t>::max()) {
         throw QueryRuntimeException("Port number invalid!");
       }
+    };
+
+    if (replication_role == ReplicationQuery::ReplicationRole::MAIN) {
+#ifdef MG_ENTERPRISE
+      if (!license::global_license_checker.IsEnterpriseValidFast()) {
+        throw QueryException("Trying to use enterprise feature without a valid license.");
+      }
+
+      ValidatePort(port);
 
       auto const config = memgraph::replication::ReplicationServerConfig{
           .ip_address = memgraph::replication::kDefaultReplicationServerIp,
           .port = static_cast<uint16_t>(*port),
       };
 
+      if (!handler_.SetReplicationRoleMain(config)) {
+        throw QueryRuntimeException("Couldn't set replication role to main!");
+      }
+#else
+      if (!handler_.SetReplicationRoleMain()) {
+        throw QueryRuntimeException("Couldn't set replication role to main!");
+      }
+#endif
+
+    }
+#ifdef MG_ENTERPRISE
+    else if (replication_role == ReplicationQuery::ReplicationRole::COORDINATOR) {
+      if (!license::global_license_checker.IsEnterpriseValidFast()) {
+        throw QueryException("Trying to use enterprise feature without a valid license.");
+      }
+
+      if (!handler_.SetReplicationRoleCoordinator()) {
+        throw QueryRuntimeException("Couldn't set replication role to coordinator!");
+      }
+    }
+#endif
+    else {  // replica
+      ValidatePort(port);
+
+      auto const config = memgraph::replication::ReplicationServerConfig{
+          .ip_address = memgraph::replication::kDefaultReplicationServerIp,
+          .port = static_cast<uint16_t>(*port),
+      };
       if (!handler_.SetReplicationRoleReplica(config)) {
-        throw QueryRuntimeException("Couldn't set role to replica!");
+        throw QueryRuntimeException("Couldn't set replication role to replica!");
       }
     }
   }
@@ -305,6 +338,13 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
         return ReplicationQuery::ReplicationRole::MAIN;
       case memgraph::replication::ReplicationRole::REPLICA:
         return ReplicationQuery::ReplicationRole::REPLICA;
+#ifdef MG_ENTERPRISE
+        if (!license::global_license_checker.IsEnterpriseValidFast()) {
+          throw QueryException("Trying to use enterprise feature without a valid license.");
+        }
+      case memgraph::replication::ReplicationRole::COORDINATOR:
+        return ReplicationQuery::ReplicationRole::COORDINATOR;
+#endif
     }
     throw QueryRuntimeException("Couldn't show replication role - invalid role set!");
   }
@@ -314,7 +354,6 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
                        const ReplicationQuery::SyncMode sync_mode,
                        const std::chrono::seconds replica_check_frequency) override {
     if (handler_.IsReplica()) {
-      // replica can't register another replica
       throw QueryRuntimeException("Replica can't register another replica!");
     }
 
@@ -328,7 +367,7 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
                                                          .mode = repl_mode,
                                                          .ip_address = ip,
                                                          .port = port,
-                                                         .replica_check_frequency = replica_check_frequency,
+                                                         .check_frequency = replica_check_frequency,
                                                          .ssl = std::nullopt};
       auto ret = handler_.RegisterReplica(config);
       if (ret.HasError()) {
@@ -339,12 +378,39 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
     }
   }
 
+#ifdef MG_ENTERPRISE
+  /// @throw QueryRuntimeException if an error ocurred.
+  void RegisterMain(const std::string &socket_address, const std::chrono::seconds main_check_frequency) override {
+    if (!handler_.IsCoordinator()) {
+      throw QueryRuntimeException("Only coordinator can register main instance!");
+    }
+
+    auto maybe_ip_and_port =
+        io::network::Endpoint::ParseSocketOrAddress(socket_address, memgraph::replication::kDefaultReplicationPort);
+    if (maybe_ip_and_port) {
+      auto [ip, port] = *maybe_ip_and_port;
+      auto config = replication::ReplicationClientConfig{.name = memgraph::replication::kDefaultMainName,
+                                                         .mode = std::nullopt,
+                                                         .ip_address = ip,
+                                                         .port = port,
+                                                         .check_frequency = main_check_frequency,
+                                                         .ssl = std::nullopt};
+      auto ret = handler_.RegisterMain(config);
+      if (ret.HasError()) {
+        throw QueryRuntimeException("Couldn't register main!");
+      }
+    } else {
+      throw QueryRuntimeException("Invalid socket address!");
+    }
+  }
+#endif
+
   /// @throw QueryRuntimeException if an error occurred.
   void DropReplica(std::string_view replica_name) override {
     auto const result = handler_.UnregisterReplica(replica_name);
     switch (result) {
       using enum memgraph::dbms::UnregisterReplicaResult;
-      case NOT_MAIN:
+      case IS_REPLICA:
         throw QueryRuntimeException("Replica can't unregister a replica!");
       case COULD_NOT_BE_PERSISTED:
         [[fallthrough]];
@@ -380,7 +446,11 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
       Replica replica;
       replica.name = repl_info.name;
       replica.socket_address = repl_info.endpoint.SocketAddress();
+#ifdef MG_ENTERPRISE
+      switch (repl_info.mode.value()) {  // Replica has always mode set
+#else
       switch (repl_info.mode) {
+#endif
         case memgraph::replication::ReplicationMode::SYNC:
           replica.sync_mode = ReplicationQuery::SyncMode::SYNC;
           break;
@@ -487,7 +557,10 @@ Callback HandleAuthQuery(AuthQuery *auth_query, InterpreterContext *interpreter_
           auth->GrantPrivilege(username, kPrivilegesAll
 #ifdef MG_ENTERPRISE
                                ,
-                               {{{AuthQuery::FineGrainedPrivilege::CREATE_DELETE, {query::kAsterisk}}}},
+                               {{
+                                 { AuthQuery::FineGrainedPrivilege::CREATE_DELETE,
+                                   { query::kAsterisk } }
+                               }},
                                {
                                  {
                                    {
@@ -753,6 +826,11 @@ Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &
           case ReplicationQuery::ReplicationRole::REPLICA: {
             return std::vector<std::vector<TypedValue>>{{TypedValue("replica")}};
           }
+#ifdef MG_ENTERPRISE
+          case ReplicationQuery::ReplicationRole::COORDINATOR: {
+            return std::vector<std::vector<TypedValue>>{{ TypedValue("coordinator") }};
+          }
+#endif
         }
       };
       return callback;
@@ -761,10 +839,9 @@ Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &
       const auto &name = repl_query->replica_name_;
       const auto &sync_mode = repl_query->sync_mode_;
       auto socket_address = repl_query->socket_address_->Accept(evaluator);
-      const auto replica_check_frequency = config.replication_replica_check_frequency;
 
       callback.fn = [handler = ReplQueryHandler{dbms_handler}, name, socket_address, sync_mode,
-                     replica_check_frequency]() mutable {
+                     replica_check_frequency = config.replication_replica_check_frequency]() mutable {
         handler.RegisterReplica(name, std::string(socket_address.ValueString()), sync_mode, replica_check_frequency);
         return std::vector<std::vector<TypedValue>>();
       };
@@ -772,6 +849,28 @@ Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &
                                   fmt::format("Replica {} is registered.", repl_query->replica_name_));
       return callback;
     }
+    case ReplicationQuery::Action::REGISTER_MAIN: {
+      if (!license::global_license_checker.IsEnterpriseValidFast()) {
+        throw QueryException("Trying to use enterprise feature without a valid license.");
+      }
+#ifdef MG_ENTERPRISE
+      // TODO: MemoryResource for EvaluationContext, it should probably be passed as
+      // the argument to Callback.
+      EvaluationContext evaluation_context{.timestamp = QueryTimestamp(), .parameters = parameters};
+      auto evaluator = PrimitiveLiteralExpressionEvaluator{evaluation_context};
+      auto socket_address_tv = repl_query->socket_address_->Accept(evaluator);
+
+      callback.fn = [handler = ReplQueryHandler{dbms_handler}, socket_address_tv,
+                     main_check_frequency = config.replication_replica_check_frequency]() mutable {
+        handler.RegisterMain(std::string(socket_address_tv.ValueString()), main_check_frequency);
+        return std::vector<std::vector<TypedValue>>();
+      };
+
+      notifications->emplace_back(SeverityLevel::INFO, NotificationCode::REGISTER_MAIN,
+                                  "Coordinator has registered main instance.");
+      return callback;
+#endif
+    }
     case ReplicationQuery::Action::DROP_REPLICA: {
       const auto &name = repl_query->replica_name_;
       callback.fn = [handler = ReplQueryHandler{dbms_handler}, name]() mutable {
@@ -2237,8 +2336,12 @@ PreparedQuery PrepareAuthQuery(ParsedQuery parsed_query, bool in_explicit_transa
   auto callback = HandleAuthQuery(auth_query, interpreter_context, parsed_query.parameters);
 
   return PreparedQuery{std::move(callback.header), std::move(parsed_query.required_privileges),
-                       [handler = std::move(callback.fn), pull_plan = std::shared_ptr<PullPlanVector>(nullptr),
-                        interpreter_context](  // NOLINT
+                       [handler = std::move(callback.fn), pull_plan = std::shared_ptr<PullPlanVector>(nullptr)
+#ifdef MG_ENTERPRISE
+                                                              ,
+                        interpreter_context
+#endif
+  ](  // NOLINT
                            AnyStream *stream, std::optional<int> n) mutable -> std::optional<QueryHandlerResult> {
                          if (!pull_plan) {
                            // Run the specific query

src/query/interpreter.hpp

@@ -99,6 +99,11 @@ class ReplicationQueryHandler {
                                ReplicationQuery::SyncMode sync_mode,
                                const std::chrono::seconds replica_check_frequency) = 0;
 
+#ifdef MG_ENTERPRISE
+  /// @throw QueryRuntimeException if an error ocurred.
+  virtual void RegisterMain(const std::string &socket_address, const std::chrono::seconds main_check_frequency) = 0;
+#endif
+
   /// @throw QueryRuntimeException if an error ocurred.
   virtual void DropReplica(std::string_view replica_name) = 0;
 

src/query/metadata.cpp

@@ -66,6 +66,10 @@ constexpr std::string_view GetCodeString(const NotificationCode code) {
       return "PlanHinting"sv;
     case NotificationCode::REGISTER_REPLICA:
       return "RegisterReplica"sv;
+#ifdef MG_ENTERPRISE
+    case NotificationCode::REGISTER_MAIN:
+      return "RegisterMain"sv;
+#endif
     case NotificationCode::REPLICA_PORT_WARNING:
       return "ReplicaPortWarning"sv;
     case NotificationCode::SET_REPLICA:

src/query/metadata.hpp

@@ -42,6 +42,9 @@ enum class NotificationCode : uint8_t {
   PLAN_HINTING,
   REPLICA_PORT_WARNING,
   REGISTER_REPLICA,
+#ifdef MG_ENTERPRISE
+  REGISTER_MAIN,
+#endif
   SET_REPLICA,
   START_STREAM,
   START_ALL_STREAMS,

src/replication/include/replication/config.hpp

@@ -22,16 +22,25 @@ namespace memgraph::replication {
 inline constexpr uint16_t kDefaultReplicationPort = 10000;
 inline constexpr auto *kDefaultReplicationServerIp = "0.0.0.0";
 
+#ifdef MG_ENTERPRISE
+// Default name which coordinator uses to distinguish main's ReplicationClient from replicas'.
+inline constexpr auto *kDefaultMainName = "main";
+#endif
+
 struct ReplicationClientConfig {
   std::string name;
+#ifdef MG_ENTERPRISE
+  std::optional<ReplicationMode> mode;
+#else
   ReplicationMode mode{};
+#endif
   std::string ip_address;
   uint16_t port{};
 
-  // The default delay between main checking/pinging replicas is 1s because
-  // that seems like a reasonable timeframe in which main should notice a
-  // replica is down.
-  std::chrono::seconds replica_check_frequency{1};
+  // The default delay between coordinator/main checking/pinging main and replicas/replicas is 1s because
+  // that seems like a reasonable timeframe in which coordinator/main should notice a
+  // main or replica is down.
+  std::chrono::seconds check_frequency{1};
 
   struct SSL {
     std::string key_file;

src/replication/include/replication/mode.hpp

@@ -12,5 +12,7 @@
 #pragma once
 #include <cstdint>
 namespace memgraph::replication {
+
 enum class ReplicationMode : std::uint8_t { SYNC, ASYNC };
-}
+
+}  // namespace memgraph::replication

src/replication/include/replication/replication_client.hpp

@@ -25,6 +25,7 @@ namespace memgraph::replication {
 template <typename F>
 concept InvocableWithStringView = std::invocable<F, std::string_view>;
 
+/// TODO: (andi) Consider adding some var type which would distinguish between MAIN and REPLICA frequency_checker
 struct ReplicationClient {
   explicit ReplicationClient(const memgraph::replication::ReplicationClientConfig &config);
 
@@ -37,8 +38,8 @@ struct ReplicationClient {
   template <InvocableWithStringView F>
   void StartFrequentCheck(F &&callback) {
     // Help the user to get the most accurate replica state possible.
-    if (replica_check_frequency_ > std::chrono::seconds(0)) {
-      replica_checker_.Run("Replica Checker", replica_check_frequency_, [this, cb = std::forward<F>(callback)] {
+    if (check_frequency_ > std::chrono::seconds(0)) {
+      replica_checker_.Run("Replica Checker", check_frequency_, [this, cb = std::forward<F>(callback)] {
         try {
           bool success = false;
           {
@@ -58,9 +59,13 @@ struct ReplicationClient {
   std::string name_;
   communication::ClientContext rpc_context_;
   rpc::Client rpc_client_;
-  std::chrono::seconds replica_check_frequency_;
+  std::chrono::seconds check_frequency_;
 
+#ifdef MG_ENTERPRISE
+  std::optional<memgraph::replication::ReplicationMode> mode_{memgraph::replication::ReplicationMode::SYNC};
+#else
   memgraph::replication::ReplicationMode mode_{memgraph::replication::ReplicationMode::SYNC};
+#endif
   // This thread pool is used for background tasks so we don't
   // block the main storage thread
   // We use only 1 thread for 2 reasons:

src/replication/include/replication/role.hpp

@@ -14,5 +14,10 @@
 #include <cstdint>
 namespace memgraph::replication {
 
+#ifdef MG_ENTERPRISE
+enum class ReplicationRole : uint8_t { MAIN, REPLICA, COORDINATOR };
+#else
 enum class ReplicationRole : uint8_t { MAIN, REPLICA };
-}
+#endif
+
+}  // namespace memgraph::replication

src/replication/include/replication/state.hpp

@@ -32,27 +32,123 @@ namespace memgraph::replication {
 
 enum class RolePersisted : uint8_t { UNKNOWN_OR_NO, YES };
 
-enum class RegisterReplicaError : uint8_t { NAME_EXISTS, END_POINT_EXISTS, COULD_NOT_BE_PERSISTED, NOT_MAIN, SUCCESS };
+// TODO: (andi) Dual definition of same error.
+
+enum class RegisterReplicaError : uint8_t {
+  NAME_EXISTS,
+  END_POINT_EXISTS,
+  COULD_NOT_BE_PERSISTED,
+  IS_REPLICA,
+  SUCCESS
+};
+
+#ifdef MG_ENTERPRISE
+enum class RegisterMainError : uint8_t {
+  MAIN_ALREADY_EXISTS,
+  END_POINT_EXISTS,
+  COULD_NOT_BE_PERSISTED,
+  NOT_COORDINATOR,
+  SUCCESS
+};
+#endif
 
 struct RoleMainData {
+  // TODO: (andi) Currently, RoleMainData can exist without server and server_config_ in enterprise version of the code.
+  // Introducing new role should solve this non-happy design decision.
   RoleMainData() = default;
-  explicit RoleMainData(ReplicationEpoch e) : epoch_(std::move(e)) {}
+
+#ifdef MG_ENTERPRISE
+  RoleMainData(ReplicationEpoch epoch, ReplicationServerConfig server_config)
+      : epoch_(std::move(epoch)),
+        server_config_(std::move(server_config)),
+        server_(std::make_unique<ReplicationServer>(server_config_)) {}
+#else
+  explicit RoleMainData(ReplicationEpoch epoch) : epoch_(std::move(epoch)) {}
+#endif
+
   ~RoleMainData() = default;
 
   RoleMainData(RoleMainData const &) = delete;
   RoleMainData &operator=(RoleMainData const &) = delete;
-  RoleMainData(RoleMainData &&) = default;
-  RoleMainData &operator=(RoleMainData &&) = default;
+
+#ifdef MG_ENTERPRISE
+  RoleMainData(RoleMainData &&other) noexcept
+      : epoch_(std::move(other.epoch_)),
+        registered_replicas_(std::move(other.registered_replicas_)),
+        server_config_(std::move(other.server_config_)),
+        server_(std::move(other.server_)) {}
+
+  RoleMainData &operator=(RoleMainData &&other) noexcept {
+    if (this != &other) {
+      epoch_ = std::move(other.epoch_);
+      registered_replicas_ = std::move(other.registered_replicas_);
+      server_config_ = std::move(other.server_config_);
+      server_ = std::move(other.server_);
+    }
+    return *this;
+  }
+#else
+  RoleMainData(RoleMainData &&) noexcept = default;
+  RoleMainData &operator=(RoleMainData &&) noexcept = default;
+#endif
 
   ReplicationEpoch epoch_;
-  std::list<ReplicationClient> registered_replicas_{};
+  std::list<ReplicationClient> registered_replicas_;
+
+#ifdef MG_ENTERPRISE
+  ReplicationServerConfig server_config_;
+  std::unique_ptr<ReplicationServer> server_;
+#endif
 };
 
 struct RoleReplicaData {
-  ReplicationServerConfig config;
-  std::unique_ptr<ReplicationServer> server;
+  explicit RoleReplicaData(ReplicationServerConfig config)
+      : config_(std::move(config)), server_(std::make_unique<ReplicationServer>(config_)) {}
+
+  ~RoleReplicaData() = default;
+
+  RoleReplicaData(RoleReplicaData const &) = delete;
+  RoleReplicaData &operator=(RoleReplicaData const &) = delete;
+
+  RoleReplicaData(RoleReplicaData &&other) noexcept
+      : config_(std::move(other.config_)), server_(std::move(other.server_)) {}
+
+  RoleReplicaData &operator=(RoleReplicaData &&other) noexcept {
+    if (this != &other) {
+      config_ = std::move(other.config_);
+      server_ = std::move(other.server_);
+    }
+    return *this;
+  }
+
+  ReplicationServerConfig config_;
+  std::unique_ptr<ReplicationServer> server_;
 };
 
+#ifdef MG_ENTERPRISE
+struct RoleCoordinatorData {
+  RoleCoordinatorData() = default;
+  ~RoleCoordinatorData() = default;
+
+  RoleCoordinatorData(RoleCoordinatorData const &) = delete;
+  RoleCoordinatorData &operator=(RoleCoordinatorData const &) = delete;
+
+  RoleCoordinatorData(RoleCoordinatorData &&other) noexcept
+      : registered_replicas_(std::move(other.registered_replicas_)), main(std::move(other.main)) {}
+  RoleCoordinatorData &operator=(RoleCoordinatorData &&other) noexcept {
+    if (this != &other) {
+      registered_replicas_ = std::move(other.registered_replicas_);
+      main = std::move(other.main);
+    }
+    return *this;
+  }
+
+  // TODO: (andi) Does it need epoch or some other way or tracking what is going on?
+  std::list<ReplicationClient> registered_replicas_;
+  std::unique_ptr<ReplicationClient> main;
+};
+#endif
+
 // Global (instance) level object
 struct ReplicationState {
   explicit ReplicationState(std::optional<std::filesystem::path> durability_dir);
@@ -68,29 +164,66 @@ struct ReplicationState {
     PARSE_ERROR,
   };
 
+#ifdef MG_ENTERPRISE
+  using ReplicationData_t = std::variant<RoleMainData, RoleReplicaData, RoleCoordinatorData>;
+#else
   using ReplicationData_t = std::variant<RoleMainData, RoleReplicaData>;
+#endif
+
   using FetchReplicationResult_t = utils::BasicResult<FetchReplicationError, ReplicationData_t>;
   auto FetchReplicationData() -> FetchReplicationResult_t;
 
   auto GetRole() const -> ReplicationRole {
-    return std::holds_alternative<RoleReplicaData>(replication_data_) ? ReplicationRole::REPLICA
-                                                                      : ReplicationRole::MAIN;
+    if (std::holds_alternative<RoleReplicaData>(replication_data_)) {
+      return ReplicationRole::REPLICA;
+    }
+#ifdef MG_ENTERPRISE
+    if (std::holds_alternative<RoleCoordinatorData>(replication_data_)) {
+      return ReplicationRole::COORDINATOR;
+    }
+#endif
+    return ReplicationRole::MAIN;
   }
+
+#ifdef MG_ENTERPRISE
+  // TODO: (andi) Unregistering main from coordinator
+  bool IsCoordinator() const { return GetRole() == ReplicationRole::COORDINATOR; }
+  bool TryPersistRoleCoordinator();
+  bool TryPersistRegisteredReplicaOnCoordinator(const ReplicationClientConfig &config);
+  bool TryPersistUnregisterReplicaOnCoordinator(std::string_view name);
+  bool TryPersistRegisteredMainOnCoordinator(const ReplicationClientConfig &config);
+#endif
+
   bool IsMain() const { return GetRole() == ReplicationRole::MAIN; }
   bool IsReplica() const { return GetRole() == ReplicationRole::REPLICA; }
-
   bool ShouldPersist() const { return nullptr != durability_; }
-  bool TryPersistRoleMain(std::string new_epoch);
+
   bool TryPersistRoleReplica(const ReplicationServerConfig &config);
-  bool TryPersistUnregisterReplica(std::string_view name);
-  bool TryPersistRegisteredReplica(const ReplicationClientConfig &config);
+  bool TryPersistRegisteredReplicaOnMain(const ReplicationClientConfig &config);
+  bool TryPersistUnregisterReplicaOnMain(std::string_view name);
+
+#ifdef MG_ENTERPRISE
+  bool TryPersistRoleMain(std::string new_epoch, const ReplicationServerConfig &config);
+#else
+  bool TryPersistRoleMain(std::string new_epoch);
+#endif
 
   // TODO: locked access
   auto ReplicationData() -> ReplicationData_t & { return replication_data_; }
   auto ReplicationData() const -> ReplicationData_t const & { return replication_data_; }
+
   utils::BasicResult<RegisterReplicaError, ReplicationClient *> RegisterReplica(const ReplicationClientConfig &config);
 
+#ifdef MG_ENTERPRISE
+  utils::BasicResult<RegisterMainError, ReplicationClient *> RegisterMain(const ReplicationClientConfig &config);
+  bool SetReplicationRoleCoordinator();
+#endif
+
+#ifdef MG_ENTERPRISE
+  bool SetReplicationRoleMain(const ReplicationServerConfig &config);
+#else
   bool SetReplicationRoleMain();
+#endif
 
   bool SetReplicationRoleReplica(const ReplicationServerConfig &config);
 

src/replication/include/replication/status.hpp

@@ -29,6 +29,10 @@ namespace memgraph::replication::durability {
 constexpr auto *kReplicationRoleName{"__replication_role"};
 constexpr auto *kReplicationReplicaPrefix{"__replication_replica:"};  // introduced in V2
 
+#ifdef MG_ENTERPRISE
+constexpr auto *kReplicationMainPrefix("__replication_main:");
+#endif
+
 enum class DurabilityVersion : uint8_t {
   V1,  // no distinct key for replicas
   V2,  // this version, epoch, replica prefix introduced
@@ -37,34 +41,49 @@ enum class DurabilityVersion : uint8_t {
 // fragment of key: "__replication_role"
 struct MainRole {
   ReplicationEpoch epoch{};
+#ifdef MG_ENTERPRISE
+  ReplicationServerConfig config{};
+#endif
   friend bool operator==(MainRole const &, MainRole const &) = default;
 };
 
 // fragment of key: "__replication_role"
 struct ReplicaRole {
-  ReplicationServerConfig config;
+  ReplicationServerConfig config{};
   friend bool operator==(ReplicaRole const &, ReplicaRole const &) = default;
 };
 
+#ifdef MG_ENTERPRISE
+// fragment of key: "__replication_role"
+// TODO: (andi) Check if you will need to add some epoch or something for managing coordinator role
+struct CoordinatorRole {
+  friend bool operator==(CoordinatorRole const &, CoordinatorRole const &) = default;
+};
+#endif
+
 // from key: "__replication_role"
 struct ReplicationRoleEntry {
   DurabilityVersion version =
       DurabilityVersion::V2;  // if not latest then migration required for kReplicationReplicaPrefix
+
+#ifdef MG_ENTERPRISE
+  std::variant<MainRole, ReplicaRole, CoordinatorRole> role;
+#else
   std::variant<MainRole, ReplicaRole> role;
+#endif
 
   friend bool operator==(ReplicationRoleEntry const &, ReplicationRoleEntry const &) = default;
 };
 
-// from key: "__replication_replica:"
-struct ReplicationReplicaEntry {
+// used for main's and replicas' clients
+struct ReplicationClientConfigEntry {
   ReplicationClientConfig config;
-  friend bool operator==(ReplicationReplicaEntry const &, ReplicationReplicaEntry const &) = default;
+  friend bool operator==(ReplicationClientConfigEntry const &, ReplicationClientConfigEntry const &) = default;
 };
 
 void to_json(nlohmann::json &j, const ReplicationRoleEntry &p);
 void from_json(const nlohmann::json &j, ReplicationRoleEntry &p);
 
-void to_json(nlohmann::json &j, const ReplicationReplicaEntry &p);
-void from_json(const nlohmann::json &j, ReplicationReplicaEntry &p);
-
+void to_json(nlohmann::json &j, const ReplicationClientConfigEntry &p);
+void from_json(const nlohmann::json &j, ReplicationClientConfigEntry &p);
 }  // namespace memgraph::replication::durability

src/replication/replication_client.cpp

@@ -24,7 +24,7 @@ ReplicationClient::ReplicationClient(const memgraph::replication::ReplicationCli
       rpc_context_{CreateClientContext(config)},
       rpc_client_{io::network::Endpoint(io::network::Endpoint::needs_resolving, config.ip_address, config.port),
                   &rpc_context_},
-      replica_check_frequency_{config.replica_check_frequency},
+      check_frequency_{config.check_frequency},
       mode_{config.mode} {}
 
 ReplicationClient::~ReplicationClient() {

src/replication/state.cpp

@@ -28,6 +28,14 @@ auto BuildReplicaKey(std::string_view name) -> std::string {
   return key;
 }
 
+#ifdef MG_ENTERPRISE
+auto BuildMainKey(std::string_view name) -> std::string {
+  auto key = std::string{durability::kReplicationMainPrefix};
+  key.append(name);
+  return key;
+}
+#endif
+
 ReplicationState::ReplicationState(std::optional<std::filesystem::path> durability_dir) {
   if (!durability_dir) return;
   auto repl_dir = *std::move(durability_dir);
@@ -77,11 +85,21 @@ bool ReplicationState::TryPersistRoleReplica(const ReplicationServerConfig &conf
   return true;
 }
 
+#ifdef MG_ENTERPRISE
+bool ReplicationState::TryPersistRoleMain(std::string new_epoch, const ReplicationServerConfig &config) {
+#else
 bool ReplicationState::TryPersistRoleMain(std::string new_epoch) {
+#endif
+
   if (!ShouldPersist()) return true;
 
+#ifdef MG_ENTERPRISE
+  auto data = durability::ReplicationRoleEntry{
+      .role = durability::MainRole{.epoch = ReplicationEpoch{std::move(new_epoch)}, .config = config}};
+#else
   auto data =
       durability::ReplicationRoleEntry{.role = durability::MainRole{.epoch = ReplicationEpoch{std::move(new_epoch)}}};
+#endif
 
   if (durability_->Put(durability::kReplicationRoleName, nlohmann::json(data).dump())) {
     role_persisted = RolePersisted::YES;
@@ -91,7 +109,23 @@ bool ReplicationState::TryPersistRoleMain(std::string new_epoch) {
   return false;
 }
 
-bool ReplicationState::TryPersistUnregisterReplica(std::string_view name) {
+#ifdef MG_ENTERPRISE
+bool ReplicationState::TryPersistRoleCoordinator() {
+  if (!ShouldPersist()) return true;
+
+  /// TODO: (andi) Epoch/identifier possible change
+  auto data = durability::ReplicationRoleEntry{.role = durability::CoordinatorRole{}};
+
+  if (durability_->Put(durability::kReplicationRoleName, nlohmann::json(data).dump())) {
+    role_persisted = RolePersisted::YES;
+    return true;
+  }
+  spdlog::error("Error when saving COORDINATOR replication role in settings.");
+  return false;
+}
+#endif
+
+bool ReplicationState::TryPersistUnregisterReplicaOnMain(std::string_view name) {
   if (!ShouldPersist()) return true;
 
   auto key = BuildReplicaKey(name);
@@ -101,6 +135,18 @@ bool ReplicationState::TryPersistUnregisterReplica(std::string_view name) {
   return false;
 }
 
+#ifdef MG_ENTERPRISE
+bool ReplicationState::TryPersistUnregisterReplicaOnCoordinator(std::string_view name) {
+  if (!ShouldPersist()) return true;
+
+  auto key = BuildReplicaKey(name);
+
+  if (durability_->Delete(key)) return true;
+  spdlog::error("Error when removing replica {} from settings.", name);
+  return false;
+}
+#endif
+
 // TODO: FetchEpochData (agnostic of FetchReplicationData, but should be done before)
 
 auto ReplicationState::FetchReplicationData() -> FetchReplicationResult_t {
@@ -121,38 +167,53 @@ auto ReplicationState::FetchReplicationData() -> FetchReplicationResult_t {
       return FetchReplicationError::PARSE_ERROR;
     }
 
-    // To get here this must be the case
     role_persisted = memgraph::replication::RolePersisted::YES;
 
-    return std::visit(
-        utils::Overloaded{
-            [&](durability::MainRole &&r) -> FetchReplicationResult_t {
-              auto res = RoleMainData{std::move(r.epoch)};
-              auto b = durability_->begin(durability::kReplicationReplicaPrefix);
-              auto e = durability_->end(durability::kReplicationReplicaPrefix);
-              for (; b != e; ++b) {
-                auto const &[replica_name, replica_data] = *b;
-                auto json = nlohmann::json::parse(replica_data, nullptr, false);
-                if (json.is_discarded()) return FetchReplicationError::PARSE_ERROR;
-                try {
-                  durability::ReplicationReplicaEntry data = json.get<durability::ReplicationReplicaEntry>();
-                  auto key_name = std::string_view{replica_name}.substr(strlen(durability::kReplicationReplicaPrefix));
-                  if (key_name != data.config.name) {
-                    return FetchReplicationError::PARSE_ERROR;
-                  }
-                  // Instance clients
-                  res.registered_replicas_.emplace_back(data.config);
-                } catch (...) {
-                  return FetchReplicationError::PARSE_ERROR;
-                }
-              }
-              return {std::move(res)};
-            },
-            [&](durability::ReplicaRole &&r) -> FetchReplicationResult_t {
-              return {RoleReplicaData{r.config, std::make_unique<ReplicationServer>(r.config)}};
-            },
-        },
-        std::move(data.role));
+    auto const main_handler = [&](durability::MainRole &&r) -> FetchReplicationResult_t {
+#ifdef MG_ENTERPRISE
+      auto res = RoleMainData{std::move(r.epoch), std::move(r.config)};
+#else
+      auto res = RoleMainData{std::move(r.epoch)};
+#endif
+      auto b = durability_->begin(durability::kReplicationReplicaPrefix);
+      auto e = durability_->end(durability::kReplicationReplicaPrefix);
+      for (; b != e; ++b) {
+        auto const &[replica_name, replica_data] = *b;
+        auto json = nlohmann::json::parse(replica_data, nullptr, false);
+        if (json.is_discarded()) {
+          return FetchReplicationError::PARSE_ERROR;
+        }
+        try {
+          durability::ReplicationClientConfigEntry data = json.get<durability::ReplicationClientConfigEntry>();
+          auto key_name = std::string_view{replica_name}.substr(strlen(durability::kReplicationReplicaPrefix));
+          if (key_name != data.config.name) {
+            return FetchReplicationError::PARSE_ERROR;
+          }
+          // Instance clients
+          res.registered_replicas_.emplace_back(data.config);
+        } catch (...) {
+          return FetchReplicationError::PARSE_ERROR;
+        }
+      }
+      return {std::move(res)};
+    };
+
+    auto const replica_handler = [&](durability::ReplicaRole &&r) -> FetchReplicationResult_t {
+      return {RoleReplicaData{r.config}};
+    };
+
+#ifdef MG_ENTERPRISE
+    auto const coordinator_handler = [&](durability::CoordinatorRole &&) -> FetchReplicationResult_t {
+      // When coordinator wakes up, we don't restore its state from before. This step will be handled once we include
+      // high availability for coordinator.
+      return {RoleCoordinatorData{}};
+    };
+
+    return std::visit(utils::Overloaded{main_handler, replica_handler, coordinator_handler}, std::move(data.role));
+#else
+    return std::visit(utils::Overloaded{main_handler, replica_handler}, std::move(data.role));
+#endif
+
   } catch (...) {
     return FetchReplicationError::PARSE_ERROR;
   }
@@ -174,12 +235,12 @@ bool ReplicationState::HandleVersionMigration(durability::ReplicationRoleEntry &
         auto old_json = nlohmann::json::parse(old_data, nullptr, false);
         if (old_json.is_discarded()) return false;  // Can not read old_data as json
         try {
-          durability::ReplicationReplicaEntry new_data = old_json.get<durability::ReplicationReplicaEntry>();
+          durability::ReplicationClientConfigEntry new_data = old_json.get<durability::ReplicationClientConfigEntry>();
 
           // Migrate to using new key
           to_put.emplace(BuildReplicaKey(old_key), nlohmann::json(new_data).dump());
         } catch (...) {
-          return false;  // Can not parse as ReplicationReplicaEntry
+          return false;  // Can not parse as ReplicationClientConfigEntry
         }
         to_delete.push_back(std::move(old_key));
       }
@@ -198,17 +259,23 @@ bool ReplicationState::HandleVersionMigration(durability::ReplicationRoleEntry &
   return true;
 }
 
-bool ReplicationState::TryPersistRegisteredReplica(const ReplicationClientConfig &config) {
+bool ReplicationState::TryPersistRegisteredReplicaOnMain(const ReplicationClientConfig &config) {
   if (!ShouldPersist()) return true;
+  DMG_ASSERT(IsMain(), "Expected MAIN when persisting registered replica.");
 
   // If any replicas are persisted then Role must be persisted
   if (role_persisted != RolePersisted::YES) {
-    DMG_ASSERT(IsMain(), "MAIN is expected");
     auto epoch_str = std::string(std::get<RoleMainData>(replication_data_).epoch_.id());
+
+#ifdef MG_ENTERPRISE
+    auto server_config = std::get<RoleMainData>(replication_data_).server_config_;
+    if (!TryPersistRoleMain(std::move(epoch_str), server_config)) return false;
+#else
     if (!TryPersistRoleMain(std::move(epoch_str))) return false;
+#endif
   }
 
-  auto data = durability::ReplicationReplicaEntry{.config = config};
+  auto data = durability::ReplicationClientConfigEntry{.config = config};
 
   auto key = BuildReplicaKey(config.name);
   if (durability_->Put(key, nlohmann::json(data).dump())) return true;
@@ -216,12 +283,63 @@ bool ReplicationState::TryPersistRegisteredReplica(const ReplicationClientConfig
   return false;
 }
 
+#ifdef MG_ENTERPRISE
+bool ReplicationState::TryPersistRegisteredReplicaOnCoordinator(const ReplicationClientConfig &config) {
+  if (!ShouldPersist()) return true;
+  MG_ASSERT(IsCoordinator(), "Expected COORDINATOR when registering replica.");
+
+  // If any replicas are persisted then Role must be persisted
+  if (role_persisted != RolePersisted::YES) {
+    if (!TryPersistRoleCoordinator()) return false;
+  }
+
+  // TODO: (andi) This is violation of DRY, try to fix it -> same code as when main registers replica.
+  auto data = durability::ReplicationClientConfigEntry{.config = config};
+  auto key = BuildReplicaKey(config.name);
+  if (durability_->Put(key, nlohmann::json(data).dump())) return true;
+  spdlog::error("Error when saving replica {} in settings.", config.name);
+  return false;
+}
+#endif
+
+#ifdef MG_ENTERPRISE
+bool ReplicationState::TryPersistRegisteredMainOnCoordinator(const ReplicationClientConfig &config) {
+  if (!ShouldPersist()) return true;
+  MG_ASSERT(IsCoordinator(), "Expected COORDINATOR when registering replica.");
+
+  // If MAIN is persisted then role must be persisted
+  if (role_persisted != RolePersisted::YES) {
+    if (!TryPersistRoleCoordinator()) return false;
+  }
+
+  auto data = durability::ReplicationClientConfigEntry{.config = config};
+  auto key = BuildMainKey(config.name);
+  if (durability_->Put(key, nlohmann::json(data).dump())) return true;
+  spdlog::error("Error when saving main {} in settings.", config.name);
+  return false;
+}
+#endif
+
+#ifdef MG_ENTERPRISE
+bool ReplicationState::SetReplicationRoleMain(const ReplicationServerConfig &config) {
+#else
 bool ReplicationState::SetReplicationRoleMain() {
+#endif
+
   auto new_epoch = utils::GenerateUUID();
+
+#ifdef MG_ENTERPRISE
+  if (!TryPersistRoleMain(new_epoch, config)) {
+    return false;
+  }
+  replication_data_ = RoleMainData{ReplicationEpoch{new_epoch}, config};
+#else
   if (!TryPersistRoleMain(new_epoch)) {
     return false;
   }
   replication_data_ = RoleMainData{ReplicationEpoch{new_epoch}};
+#endif
+
   return true;
 }
 
@@ -229,38 +347,54 @@ bool ReplicationState::SetReplicationRoleReplica(const ReplicationServerConfig &
   if (!TryPersistRoleReplica(config)) {
     return false;
   }
-  replication_data_ = RoleReplicaData{config, std::make_unique<ReplicationServer>(config)};
+  replication_data_ = RoleReplicaData{config};
   return true;
 }
 
+#ifdef MG_ENTERPRISE
+bool ReplicationState::SetReplicationRoleCoordinator() {
+  /// TODO: (andi) Think if need epoch and how is this going to be tracked.
+  if (!TryPersistRoleCoordinator()) {
+    return false;
+  }
+
+  replication_data_ = RoleCoordinatorData{};
+  return true;
+}
+#endif
+
 utils::BasicResult<RegisterReplicaError, ReplicationClient *> ReplicationState::RegisterReplica(
     const ReplicationClientConfig &config) {
-  auto const replica_handler = [](RoleReplicaData const &) { return RegisterReplicaError::NOT_MAIN; };
+  auto const replica_handler = [](RoleReplicaData const &) { return RegisterReplicaError::IS_REPLICA; };
+
+  // Returned for MAIN and COORDINATOR
   ReplicationClient *client{nullptr};
-  auto const main_handler = [&client, &config, this](RoleMainData &mainData) -> RegisterReplicaError {
-    // name check
-    auto name_check = [&config](auto const &replicas) {
-      auto name_matches = [&name = config.name](auto const &replica) { return replica.name_ == name; };
-      return std::any_of(replicas.begin(), replicas.end(), name_matches);
+
+  auto name_check = [&config](auto const &replicas) {
+    auto name_matches = [&name = config.name](auto const &replica) { return replica.name_ == name; };
+    return std::any_of(replicas.begin(), replicas.end(), name_matches);
+  };
+
+  auto endpoint_check = [&](auto const &replicas) {
+    auto endpoint_matches = [&config](auto const &replica) {
+      const auto &ep = replica.rpc_client_.Endpoint();
+      return ep.address == config.ip_address && ep.port == config.port;
     };
+    return std::any_of(replicas.begin(), replicas.end(), endpoint_matches);
+  };
+
+  auto const main_handler = [&client, &config, &name_check, &endpoint_check,
+                             this](RoleMainData &mainData) -> RegisterReplicaError {
     if (name_check(mainData.registered_replicas_)) {
       return RegisterReplicaError::NAME_EXISTS;
     }
 
-    // endpoint check
-    auto endpoint_check = [&](auto const &replicas) {
-      auto endpoint_matches = [&config](auto const &replica) {
-        const auto &ep = replica.rpc_client_.Endpoint();
-        return ep.address == config.ip_address && ep.port == config.port;
-      };
-      return std::any_of(replicas.begin(), replicas.end(), endpoint_matches);
-    };
     if (endpoint_check(mainData.registered_replicas_)) {
       return RegisterReplicaError::END_POINT_EXISTS;
     }
 
     // Durability
-    if (!TryPersistRegisteredReplica(config)) {
+    if (!TryPersistRegisteredReplicaOnMain(config)) {
       return RegisterReplicaError::COULD_NOT_BE_PERSISTED;
     }
 
@@ -269,10 +403,78 @@ utils::BasicResult<RegisterReplicaError, ReplicationClient *> ReplicationState::
     return RegisterReplicaError::SUCCESS;
   };
 
+#ifdef MG_ENTERPRISE
+  auto const coordinator_handler = [&client, &config, &name_check, &endpoint_check,
+                                    this](RoleCoordinatorData &coordinatorData) -> RegisterReplicaError {
+    if (name_check(coordinatorData.registered_replicas_)) {
+      return RegisterReplicaError::NAME_EXISTS;
+    }
+
+    if (endpoint_check(coordinatorData.registered_replicas_)) {
+      return RegisterReplicaError::END_POINT_EXISTS;
+    }
+
+    if (!TryPersistRegisteredReplicaOnCoordinator(config)) {
+      return RegisterReplicaError::COULD_NOT_BE_PERSISTED;
+    }
+
+    client = &coordinatorData.registered_replicas_.emplace_back(config);
+    return RegisterReplicaError::SUCCESS;
+  };
+
+  const auto &res =
+      std::visit(utils::Overloaded{main_handler, replica_handler, coordinator_handler}, replication_data_);
+#else
   const auto &res = std::visit(utils::Overloaded{main_handler, replica_handler}, replication_data_);
+#endif
+
   if (res == RegisterReplicaError::SUCCESS) {
     return client;
   }
   return res;
 }
+
+#ifdef MG_ENTERPRISE
+utils::BasicResult<RegisterMainError, ReplicationClient *> ReplicationState::RegisterMain(
+    const ReplicationClientConfig &config) {
+  // TODO: (andi) RegisterReplicaStatus not Error since you also have Success
+  auto const replica_handler = [](RoleReplicaData const &) { return RegisterMainError::NOT_COORDINATOR; };
+  auto const main_handler = [](RoleMainData const &) { return RegisterMainError::NOT_COORDINATOR; };
+
+  ReplicationClient *client{nullptr};
+  auto const coordinator_handler = [&config, &client, this](RoleCoordinatorData &coordinatorData) -> RegisterMainError {
+    auto name_matches = [](auto const &replica) { return replica.name_ == kDefaultMainName; };
+    if (std::any_of(coordinatorData.registered_replicas_.begin(), coordinatorData.registered_replicas_.end(),
+                    name_matches)) {
+      return RegisterMainError::MAIN_ALREADY_EXISTS;
+    }
+
+    auto endpoint_matches = [&config](auto const &replica) {
+      const auto &ep = replica.rpc_client_.Endpoint();
+      return ep.address == config.ip_address && ep.port == config.port;
+    };
+    if (std::any_of(coordinatorData.registered_replicas_.begin(), coordinatorData.registered_replicas_.end(),
+                    endpoint_matches)) {
+      return RegisterMainError::END_POINT_EXISTS;
+    }
+
+    if (!TryPersistRegisteredMainOnCoordinator(config)) {
+      return RegisterMainError::COULD_NOT_BE_PERSISTED;
+    }
+
+    coordinatorData.main = std::make_unique<ReplicationClient>(config);
+    client = coordinatorData.main.get();
+    return RegisterMainError::SUCCESS;
+  };
+
+  const auto &res =
+      std::visit(utils::Overloaded{main_handler, replica_handler, coordinator_handler}, replication_data_);
+
+  if (res == RegisterMainError::SUCCESS) {
+    return client;
+  }
+  return res;
+}
+#endif
+
 }  // namespace memgraph::replication

src/replication/status.cpp

@@ -28,9 +28,20 @@ constexpr auto *kEpoch = "epoch";
 constexpr auto *kVersion = "durability_version";
 
 void to_json(nlohmann::json &j, const ReplicationRoleEntry &p) {
+#ifdef MG_ENTERPRISE
+  auto processMAIN = [&](MainRole const &main) {
+    j = nlohmann::json{{kVersion, p.version},
+                       {kReplicationRole, ReplicationRole::MAIN},
+                       {kEpoch, main.epoch.id()},
+                       {kIpAddress, main.config.ip_address},
+                       {kPort, main.config.port}};
+  };
+#else
   auto processMAIN = [&](MainRole const &main) {
     j = nlohmann::json{{kVersion, p.version}, {kReplicationRole, ReplicationRole::MAIN}, {kEpoch, main.epoch.id()}};
   };
+#endif
+
   auto processREPLICA = [&](ReplicaRole const &replica) {
     j = nlohmann::json{
         {kVersion, p.version},
@@ -40,7 +51,17 @@ void to_json(nlohmann::json &j, const ReplicationRoleEntry &p) {
         // TODO: SSL
     };
   };
+
+#ifdef MG_ENTERPRISE
+  /// TODO: (andi) Extend this if you need to introduce some kind of timestamp to coordinator.
+  auto processCOORDINATOR = [&](CoordinatorRole const &) {
+    j = nlohmann::json{{kVersion, p.version}, {kReplicationRole, ReplicationRole::COORDINATOR}};
+  };
+
+  std::visit(utils::Overloaded{processMAIN, processREPLICA, processCOORDINATOR}, p.role);
+#else
   std::visit(utils::Overloaded{processMAIN, processREPLICA}, p.role);
+#endif
 }
 
 void from_json(const nlohmann::json &j, ReplicationRoleEntry &p) {
@@ -49,33 +70,57 @@ void from_json(const nlohmann::json &j, ReplicationRoleEntry &p) {
   // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
   ReplicationRole role;
   j.at(kReplicationRole).get_to(role);
+
+  auto ParseReplicationServerConfig = [](const nlohmann::json &j) -> ReplicationServerConfig {
+    std::string ip_address;
+    // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
+    uint16_t port;
+    j.at(kIpAddress).get_to(ip_address);
+    j.at(kPort).get_to(port);
+    return ReplicationServerConfig{.ip_address = std::move(ip_address), .port = port};
+  };
+
   switch (role) {
     case ReplicationRole::MAIN: {
       auto json_epoch = j.value(kEpoch, std::string{});
       auto epoch = ReplicationEpoch{};
       if (!json_epoch.empty()) epoch.SetEpoch(json_epoch);
+#ifdef MG_ENTERPRISE
+      auto config = ParseReplicationServerConfig(j);
+      p = ReplicationRoleEntry{.version = version,
+                               .role = MainRole{.epoch = std::move(epoch), .config = std::move(config)}};
+#else
       p = ReplicationRoleEntry{.version = version, .role = MainRole{.epoch = std::move(epoch)}};
+#endif
       break;
     }
     case ReplicationRole::REPLICA: {
-      std::string ip_address;
-      // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
-      uint16_t port;
-      j.at(kIpAddress).get_to(ip_address);
-      j.at(kPort).get_to(port);
-      auto config = ReplicationServerConfig{.ip_address = std::move(ip_address), .port = port};
+      auto config = ParseReplicationServerConfig(j);
       p = ReplicationRoleEntry{.version = version, .role = ReplicaRole{.config = std::move(config)}};
       break;
     }
+#ifdef MG_ENTERPRISE
+    case ReplicationRole::COORDINATOR: {
+      p = ReplicationRoleEntry{.version = version, .role = CoordinatorRole{}};
+      break;
+    }
+#endif
   }
 }
 
-void to_json(nlohmann::json &j, const ReplicationReplicaEntry &p) {
+void to_json(nlohmann::json &j, const ReplicationClientConfigEntry &p) {
+#ifdef MG_ENTERPRISE
+  auto common = nlohmann::json{{kReplicaName, p.config.name},
+                               {kIpAddress, p.config.ip_address},
+                               {kPort, p.config.port},
+                               {kCheckFrequency, p.config.check_frequency.count()}};
+#else
   auto common = nlohmann::json{{kReplicaName, p.config.name},
                                {kIpAddress, p.config.ip_address},
                                {kPort, p.config.port},
                                {kSyncMode, p.config.mode},
-                               {kCheckFrequency, p.config.replica_check_frequency.count()}};
+                               {kCheckFrequency, p.config.check_frequency.count()}};
+#endif
 
   if (p.config.ssl.has_value()) {
     common[kSSLKeyFile] = p.config.ssl->key_file;
@@ -84,27 +129,55 @@ void to_json(nlohmann::json &j, const ReplicationReplicaEntry &p) {
     common[kSSLKeyFile] = nullptr;
     common[kSSLCertFile] = nullptr;
   }
+
+#ifdef MG_ENTERPRISE
+  if (p.config.mode.has_value()) {
+    common[kSyncMode] = p.config.mode.value();
+  } else {
+    common[kSyncMode] = nullptr;
+  }
+#endif
+
   j = std::move(common);
 }
-void from_json(const nlohmann::json &j, ReplicationReplicaEntry &p) {
+void from_json(const nlohmann::json &j, ReplicationClientConfigEntry &p) {
   const auto &key_file = j.at(kSSLKeyFile);
   const auto &cert_file = j.at(kSSLCertFile);
 
   MG_ASSERT(key_file.is_null() == cert_file.is_null());
 
   auto seconds = j.at(kCheckFrequency).get<std::chrono::seconds::rep>();
+
+#ifdef MG_ENTERPRISE
+  auto config = ReplicationClientConfig{
+      .name = j.at(kReplicaName).get<std::string>(),
+      .ip_address = j.at(kIpAddress).get<std::string>(),
+      .port = j.at(kPort).get<uint16_t>(),
+      .check_frequency = std::chrono::seconds{seconds},
+  };
+#else
   auto config = ReplicationClientConfig{
       .name = j.at(kReplicaName).get<std::string>(),
       .mode = j.at(kSyncMode).get<ReplicationMode>(),
       .ip_address = j.at(kIpAddress).get<std::string>(),
       .port = j.at(kPort).get<uint16_t>(),
-      .replica_check_frequency = std::chrono::seconds{seconds},
+      .check_frequency = std::chrono::seconds{seconds},
   };
+#endif
+
   if (!key_file.is_null()) {
     config.ssl = ReplicationClientConfig::SSL{};
     key_file.get_to(config.ssl->key_file);
     cert_file.get_to(config.ssl->cert_file);
   }
-  p = ReplicationReplicaEntry{.config = std::move(config)};
+
+#ifdef MG_ENTERPRISE
+  if (const auto &sync_mode = j.at(kSyncMode); !sync_mode.is_null()) {
+    config.mode = sync_mode.get<ReplicationMode>();
+  }
+#endif
+
+  p = ReplicationClientConfigEntry{.config = std::move(config)};
 }
+
 }  // namespace memgraph::replication::durability

src/storage/v2/replication/global.hpp

@@ -31,7 +31,11 @@ struct TimestampInfo {
 
 struct ReplicaInfo {
   std::string name;
+#ifdef MG_ENTERPRISE
+  std::optional<memgraph::replication::ReplicationMode> mode;
+#else
   memgraph::replication::ReplicationMode mode;
+#endif
   io::network::Endpoint endpoint;
   replication::ReplicaState state;
   TimestampInfo timestamp_info;

src/storage/v2/replication/replication_client.hpp

@@ -93,7 +93,12 @@ class ReplicationStorageClient {
   ~ReplicationStorageClient() = default;
 
   // TODO Remove the client related functions
+#ifdef MG_ENTERPRISE
+  auto Mode() const -> std::optional<memgraph::replication::ReplicationMode> { return client_.mode_; }
+#else
   auto Mode() const -> memgraph::replication::ReplicationMode { return client_.mode_; }
+#endif
+
   auto Name() const -> std::string const & { return client_.name_; }
   auto Endpoint() const -> io::network::Endpoint const & { return client_.rpc_client_.Endpoint(); }
 

tests/e2e/disk_storage/replication_disabled.py

@@ -15,11 +15,12 @@ import pytest
 from common import connect, execute_and_fetch_all
 
 
+# TODO: (andi) Test correct error message using pytest.raises
 def test_replication_is_disabled(connect):
     cursor = connect.cursor()
     execute_and_fetch_all(cursor, "STORAGE MODE ON_DISK_TRANSACTIONAL")
     try:
-        execute_and_fetch_all(cursor, "SET REPLICATION ROLE TO MAIN")
+        execute_and_fetch_all(cursor, "SET REPLICATION ROLE TO MAIN WITH PORT 12000")
         assert False
     except:
         assert True

tests/e2e/replication/CMakeLists.txt

@@ -14,6 +14,8 @@ copy_e2e_python_files(replication_show conftest.py)
 copy_e2e_python_files(replication_show show.py)
 copy_e2e_python_files(replication_show show_while_creating_invalid_state.py)
 copy_e2e_python_files(replication_show edge_delete.py)
+copy_e2e_python_files(replication_show set_coordinator_role.py)
+copy_e2e_python_files(replication_show manual_failover.py)
 copy_e2e_python_files_from_parent_folder(replication_show ".." memgraph.py)
 copy_e2e_python_files_from_parent_folder(replication_show ".." interactive_mg_runner.py)
 copy_e2e_python_files_from_parent_folder(replication_show ".." mg_utils.py)

tests/e2e/replication/common.py

@@ -9,13 +9,12 @@
 # by the Apache License, Version 2.0, included in the file
 # licenses/APL.txt.
 
-import mgclient
 import typing
 
+import mgclient
 
-def execute_and_fetch_all(
-    cursor: mgclient.Cursor, query: str, params: dict = {}
-) -> typing.List[tuple]:
+
+def execute_and_fetch_all(cursor: mgclient.Cursor, query: str, params: dict = {}) -> typing.List[tuple]:
     cursor.execute(query, params)
     return cursor.fetchall()
 
@@ -24,3 +23,9 @@ def connect(**kwargs) -> mgclient.Connection:
     connection = mgclient.connect(**kwargs)
     connection.autocommit = True
     return connection
+
+
+def connect_default(**kwargs) -> mgclient.Connection:
+    def_connection = mgclient.connect(host="localhost", port=7687, **kwargs)
+    def_connection.autocommit = True
+    return def_connection

tests/e2e/replication/edge_delete.py

@@ -10,7 +10,6 @@
 # licenses/APL.txt.
 
 import sys
-import time
 
 import pytest
 from common import execute_and_fetch_all

tests/e2e/replication/manual_failover.py

@@ -0,0 +1,51 @@
+# Copyright 2023 Memgraph Ltd.
+#
+# Use of this software is governed by the Business Source License
+# included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+# License, and you may not use this file except in compliance with the Business Source License.
+#
+# As of the Change Date specified in that file, in accordance with
+# the Business Source License, use of this software will be governed
+# by the Apache License, Version 2.0, included in the file
+# licenses/APL.txt.
+
+import sys
+
+import pytest
+from common import execute_and_fetch_all
+
+
+@pytest.mark.parametrize("port, role", [(7687, "main"), (7688, "replica"), (7689, "replica"), (7690, "coordinator")])
+def test_replication_cluster_is_up(port, role, connection):
+    cursor = connection(port, role).cursor()
+    data = execute_and_fetch_all(cursor, "SHOW REPLICATION ROLE;")
+    assert cursor.description[0].name == "replication role"
+    assert data[0][0] == role
+
+
+@pytest.mark.parametrize("port, role", [(7688, "replica"), (7689, "replica")])
+def test_replica_cannot_register_replica(port, role, connection):
+    cursor = connection(port, role).cursor()
+    with pytest.raises(Exception) as e:
+        execute_and_fetch_all(cursor, "REGISTER REPLICA replica_1 SYNC TO '127.0.0.1:10001'")
+    assert str(e.value) == "Replica can't register another replica!"
+
+
+@pytest.mark.parametrize("port, role", [(7688, "replica"), (7689, "replica")])
+def test_replica_cannot_become_coordinator(port, role, connection):
+    cursor = connection(port, role).cursor()
+    with pytest.raises(Exception) as e:
+        execute_and_fetch_all(cursor, "SET REPLICATION ROLE TO COORDINATOR;")
+    assert str(e.value) == "Couldn't set replication role to coordinator!"
+
+
+@pytest.mark.parametrize("port, role", [(7687, "main"), (7688, "replica"), (7689, "replica")])
+def test_main_and_replica_cannot_register_main(port, role, connection):
+    cursor = connection(port, role).cursor()
+    with pytest.raises(Exception) as e:
+        execute_and_fetch_all(cursor, "REGISTER MAIN TO '127.0.0.1:10001';")
+    assert str(e.value) == "Only coordinator can register main instance!"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-rA"]))

tests/e2e/replication/set_coordinator_role.py

@@ -0,0 +1,34 @@
+# Copyright 2023 Memgraph Ltd.
+#
+# Use of this software is governed by the Business Source License
+# included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+# License, and you may not use this file except in compliance with the Business Source License.
+#
+# As of the Change Date specified in that file, in accordance with
+# the Business Source License, use of this software will be governed
+# by the Apache License, Version 2.0, included in the file
+# licenses/APL.txt.
+
+import sys
+
+import pytest
+from common import connect_default, execute_and_fetch_all
+
+
+def test_coordinator_role_port_throws():
+    cursor = connect_default().cursor()
+    with pytest.raises(Exception) as e:
+        execute_and_fetch_all(cursor, "SET REPLICATION ROLE TO COORDINATOR WITH PORT 1011;")
+    assert str(e.value) == "Port shouldn't be specified when setting replication role to coordinator!"
+
+
+def test_coordinator_role_no_port():
+    cursor = connect_default().cursor()
+    execute_and_fetch_all(cursor, "SET REPLICATION ROLE TO COORDINATOR;")
+    res = execute_and_fetch_all(cursor, "SHOW REPLICATION ROLE;")
+    assert cursor.description[0].name == "replication role"
+    assert res[0][0] == "coordinator"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-rA"]))

tests/e2e/replication/show.py

@@ -12,8 +12,6 @@
 import sys
 
 import pytest
-import time
-
 from common import execute_and_fetch_all
 from mg_utils import mg_sleep_and_assert
 

tests/e2e/replication/workloads.yaml

@@ -52,7 +52,53 @@ template_cluster: &template_cluster
       ]
       <<: *template_validation_queries
 
+set_coordinator_role_cluster: &set_coordinator_role_cluster
+  cluster:
+    main:
+      args: ["--bolt-port", "7687", "--log-level=TRACE"]
+      log_file: "set_coordinator_role_e2e.log"
+      setup_queries: []
+      validation_queries: []
+
+manual_failover_cluster: &manual_failover_cluster
+  cluster:
+    replica_1:
+      args: [ "--bolt-port", "7688", "--log-level=TRACE" ]
+      log_file: "replication-e2e-replica1.log"
+      setup_queries: [ "SET REPLICATION ROLE TO REPLICA WITH PORT 10001;" ]
+    replica_2:
+      args: ["--bolt-port", "7689", "--log-level=TRACE"]
+      log_file: "replication-e2e-replica2.log"
+      setup_queries: ["SET REPLICATION ROLE TO REPLICA WITH PORT 10002;"]
+    main:
+      args: ["--bolt-port", "7687", "--log-level=TRACE"]
+      log_file: "replication-e2e-main.log"
+      setup_queries: [
+        "SET REPLICATION ROLE TO MAIN WITH PORT 10003;",
+        "REGISTER REPLICA replica_1 SYNC TO '127.0.0.1:10001'",
+        "REGISTER REPLICA replica_2 SYNC TO '127.0.0.1:10002'"
+      ]
+    coordinator:
+      args: ["--bolt-port", "7690", "--log-level=TRACE"]
+      log_file: "replication-e2e-coordinator.log"
+      setup_queries: [
+        "SET REPLICATION ROLE TO COORDINATOR;",
+        "REGISTER MAIN TO '127.0.0.1:10003';",
+        "REGISTER REPLICA replica_1 SYNC TO '127.0.0.1:10001'",
+        "REGISTER REPLICA replica_2 SYNC TO '127.0.0.1:10002'"
+      ]
+
 workloads:
+  - name: "Set coordinator role"
+    binary: "tests/e2e/pytest_runner.sh"
+    args: ["replication/set_coordinator_role.py"]
+    <<: *set_coordinator_role_cluster
+
+  - name: "Manual failover"
+    binary: "tests/e2e/pytest_runner.sh"
+    args: ["replication/manual_failover.py"]
+    <<: *manual_failover_cluster
+
   - name: "Constraints"
     binary: "tests/e2e/replication/memgraph__e2e__replication__constraints"
     args: []
@@ -68,7 +114,16 @@ workloads:
     args: []
     <<: *template_cluster
 
-  - name: "Show"
+  - name: "Show while creating invalid state"
+    binary: "tests/e2e/pytest_runner.sh"
+    args: ["replication/show_while_creating_invalid_state.py"]
+
+  - name: "Delete edge replication"
+    binary: "tests/e2e/pytest_runner.sh"
+    args: ["replication/edge_delete.py"]
+    <<: *template_simple_cluster
+
+  - name: "show"
     binary: "tests/e2e/pytest_runner.sh"
     args: ["replication/show.py"]
     cluster:
@@ -96,12 +151,3 @@ workloads:
           "REGISTER REPLICA replica_3 ASYNC TO '127.0.0.1:10003'"
         ]
         validation_queries: []
-
-  - name: "Show while creating invalid state"
-    binary: "tests/e2e/pytest_runner.sh"
-    args: ["replication/show_while_creating_invalid_state.py"]
-
-  - name: "Delete edge replication"
-    binary: "tests/e2e/pytest_runner.sh"
-    args: ["replication/edge_delete.py"]
-    <<: *template_simple_cluster

tests/unit/cypher_main_visitor.cpp

@@ -2544,18 +2544,15 @@ TEST_P(CypherMainVisitorTest, ShowUsersForRole) {
 }
 
 void check_replication_query(Base *ast_generator, const ReplicationQuery *query, const std::string name,
-                             const std::optional<TypedValue> socket_address, const ReplicationQuery::SyncMode sync_mode,
-                             const std::optional<TypedValue> port = {}) {
+                             const std::optional<TypedValue> socket_address,
+                             const ReplicationQuery::SyncMode sync_mode) {
   EXPECT_EQ(query->replica_name_, name);
   EXPECT_EQ(query->sync_mode_, sync_mode);
   ASSERT_EQ(static_cast<bool>(query->socket_address_), static_cast<bool>(socket_address));
   if (socket_address) {
     ast_generator->CheckLiteral(query->socket_address_, *socket_address);
   }
-  ASSERT_EQ(static_cast<bool>(query->port_), static_cast<bool>(port));
-  if (port) {
-    ast_generator->CheckLiteral(query->port_, *port);
-  }
+  ASSERT_FALSE(static_cast<bool>(query->port_));
 }
 
 TEST_P(CypherMainVisitorTest, TestShowReplicationMode) {
@@ -2587,14 +2584,15 @@ TEST_P(CypherMainVisitorTest, TestSetReplicationMode) {
 
   {
     const std::string query = "SET REPLICATION ROLE TO MAIN";
-    auto *parsed_query = dynamic_cast<ReplicationQuery *>(ast_generator.ParseQuery(query));
-    EXPECT_EQ(parsed_query->action_, ReplicationQuery::Action::SET_REPLICATION_ROLE);
-    EXPECT_EQ(parsed_query->role_, ReplicationQuery::ReplicationRole::MAIN);
+    ASSERT_THROW(ast_generator.ParseQuery(query), SemanticException);
   }
 
   {
     const std::string query = "SET REPLICATION ROLE TO MAIN WITH PORT 10000";
-    ASSERT_THROW(ast_generator.ParseQuery(query), SemanticException);
+    auto *parsed_query = dynamic_cast<ReplicationQuery *>(ast_generator.ParseQuery(query));
+    EXPECT_EQ(parsed_query->action_, ReplicationQuery::Action::SET_REPLICATION_ROLE);
+    EXPECT_EQ(parsed_query->role_, ReplicationQuery::ReplicationRole::MAIN);
+    ast_generator.CheckLiteral(parsed_query->port_, 10000);
   }
 
   {
@@ -2606,6 +2604,7 @@ TEST_P(CypherMainVisitorTest, TestSetReplicationMode) {
   }
 }
 
+// NOTE: When using RegisterReplica query, port is not used, rather just socketAdress which also saves the port.
 TEST_P(CypherMainVisitorTest, TestRegisterReplicationQuery) {
   auto &ast_generator = *GetParam();
 
@@ -2627,6 +2626,23 @@ TEST_P(CypherMainVisitorTest, TestRegisterReplicationQuery) {
                           ReplicationQuery::SyncMode::SYNC);
 }
 
+// NOTE: When using RegisterMain query, port_ is not used, rather just socket_address_ is used.
+TEST_P(CypherMainVisitorTest, TestRegisterMainQuery) {
+  auto &ast_generator = *GetParam();
+
+  // Missing IP address and port
+  const std::string faulty_query = "REGISTER MAIN TO";
+  ASSERT_THROW(ast_generator.ParseQuery(faulty_query), SyntaxException);
+
+  // Full valid query
+  std::string full_query = R"(REGISTER MAIN TO "127.0.0.1:10003")";
+  auto *full_query_parsed = dynamic_cast<ReplicationQuery *>(ast_generator.ParseQuery(full_query));
+  ASSERT_TRUE(full_query_parsed);
+
+  ast_generator.CheckLiteral(full_query_parsed->socket_address_, "127.0.0.1:10003");
+  ASSERT_EQ(full_query_parsed->port_, nullptr);
+}
+
 TEST_P(CypherMainVisitorTest, TestDeleteReplica) {
   auto &ast_generator = *GetParam();
 

tests/unit/replication_persistence_helper.cpp

@@ -23,17 +23,18 @@
 using namespace memgraph::replication::durability;
 using namespace memgraph::replication;
 
-static_assert(sizeof(ReplicationRoleEntry) == 168,
-              "Most likely you modified ReplicationRoleEntry without updating the tests. ");
-
-static_assert(sizeof(ReplicationReplicaEntry) == 160,
-              "Most likely you modified ReplicationReplicaEntry without updating the tests.");
-
 TEST(ReplicationDurability, V1Main) {
-  auto const role_entry = ReplicationRoleEntry{.version = DurabilityVersion::V1,
-                                               .role = MainRole{
-                                                   .epoch = ReplicationEpoch{"TEST_STRING"},
-                                               }};
+  auto const role_entry =
+#ifdef MG_ENTERPRISE
+      ReplicationRoleEntry{.version = DurabilityVersion::V1,
+                           .role = MainRole{
+                               .epoch = ReplicationEpoch{"TEST_STRING"},
+                               .config = ReplicationServerConfig{.ip_address = "000.123.456.789", .port = 2023},
+                           }};
+#else
+      ReplicationRoleEntry{.version = DurabilityVersion::V1,
+                           .role = MainRole{.epoch = ReplicationEpoch{"TEST_STRING"}}};
+#endif
   nlohmann::json j;
   to_json(j, role_entry);
   ReplicationRoleEntry deser;
@@ -42,10 +43,17 @@ TEST(ReplicationDurability, V1Main) {
 }
 
 TEST(ReplicationDurability, V2Main) {
-  auto const role_entry = ReplicationRoleEntry{.version = DurabilityVersion::V2,
-                                               .role = MainRole{
-                                                   .epoch = ReplicationEpoch{"TEST_STRING"},
-                                               }};
+  auto const role_entry =
+#ifdef MG_ENTERPRISE
+      ReplicationRoleEntry{.version = DurabilityVersion::V2,
+                           .role = MainRole{
+                               .epoch = ReplicationEpoch{"TEST_STRING"},
+                               .config = ReplicationServerConfig{.ip_address = "000.123.456.789", .port = 2023},
+                           }};
+#else
+      ReplicationRoleEntry{.version = DurabilityVersion::V2,
+                           .role = MainRole{.epoch = ReplicationEpoch{"TEST_STRING"}}};
+#endif
   nlohmann::json j;
   to_json(j, role_entry);
   ReplicationRoleEntry deser;
@@ -79,36 +87,47 @@ TEST(ReplicationDurability, V2Replica) {
   ASSERT_EQ(role_entry, deser);
 }
 
-TEST(ReplicationDurability, ReplicaEntrySync) {
+#ifdef MG_ENTERPRISE
+TEST(ReplicationDurability, V2Coordinator) {
+  auto const role_entry = ReplicationRoleEntry{.version = DurabilityVersion::V2, .role = CoordinatorRole{}};
+  nlohmann::json j;
+  to_json(j, role_entry);
+  ReplicationRoleEntry deser;
+  from_json(j, deser);
+  ASSERT_EQ(role_entry, deser);
+}
+#endif
+
+TEST(ReplicationDurability, ReplicaClientConfigEntrySync) {
   using namespace std::chrono_literals;
   using namespace std::string_literals;
-  auto const replica_entry = ReplicationReplicaEntry{.config = ReplicationClientConfig{
-                                                         .name = "TEST_NAME"s,
-                                                         .mode = ReplicationMode::SYNC,
-                                                         .ip_address = "000.123.456.789"s,
-                                                         .port = 2023,
-                                                         .replica_check_frequency = 3s,
-                                                     }};
+  auto const replica_entry = ReplicationClientConfigEntry{.config = ReplicationClientConfig{
+                                                              .name = "TEST_NAME"s,
+                                                              .mode = ReplicationMode::SYNC,
+                                                              .ip_address = "000.123.456.789"s,
+                                                              .port = 2023,
+                                                              .check_frequency = 3s,
+                                                          }};
   nlohmann::json j;
   to_json(j, replica_entry);
-  ReplicationReplicaEntry deser;
+  ReplicationClientConfigEntry deser;
   from_json(j, deser);
   ASSERT_EQ(replica_entry, deser);
 }
 
-TEST(ReplicationDurability, ReplicaEntryAsync) {
+TEST(ReplicationDurability, ReplicaClientConfigEntryAsync) {
   using namespace std::chrono_literals;
   using namespace std::string_literals;
-  auto const replica_entry = ReplicationReplicaEntry{.config = ReplicationClientConfig{
-                                                         .name = "TEST_NAME"s,
-                                                         .mode = ReplicationMode::ASYNC,
-                                                         .ip_address = "000.123.456.789"s,
-                                                         .port = 2023,
-                                                         .replica_check_frequency = 3s,
-                                                     }};
+  auto const replica_entry = ReplicationClientConfigEntry{.config = ReplicationClientConfig{
+                                                              .name = "TEST_NAME"s,
+                                                              .mode = ReplicationMode::ASYNC,
+                                                              .ip_address = "000.123.456.789"s,
+                                                              .port = 2023,
+                                                              .check_frequency = 3s,
+                                                          }};
   nlohmann::json j;
   to_json(j, replica_entry);
-  ReplicationReplicaEntry deser;
+  ReplicationClientConfigEntry deser;
   from_json(j, deser);
   ASSERT_EQ(replica_entry, deser);
 }

tests/unit/storage_v2_replication.cpp

@@ -741,7 +741,10 @@ TEST_F(ReplicationTest, EpochTest) {
   main.repl_handler.UnregisterReplica(replicas[0]);
   main.repl_handler.UnregisterReplica(replicas[1]);
 
-  ASSERT_TRUE(replica1.repl_handler.SetReplicationRoleMain());
+  ASSERT_TRUE(replica1.repl_handler.SetReplicationRoleMain(ReplicationServerConfig{
+      .ip_address = local_host,
+      .port = 10003,
+  }));
 
   ASSERT_FALSE(replica1.repl_handler
                    .RegisterReplica(ReplicationClientConfig{
@@ -929,80 +932,16 @@ TEST_F(ReplicationTest, ReplicationReplicaWithExistingEndPoint) {
                   .GetError() == RegisterReplicaError::END_POINT_EXISTS);
 }
 
-TEST_F(ReplicationTest, RestoringReplicationAtStartupAfterDroppingReplica) {
-  auto main_config = main_conf;
-  auto replica1_config = main_conf;
-  auto replica2_config = main_conf;
-  main_config.durability.restore_replication_state_on_startup = true;
-
-  std::filesystem::path replica1_storage_directory{std::filesystem::temp_directory_path() / "replica1"};
-  std::filesystem::path replica2_storage_directory{std::filesystem::temp_directory_path() / "replica2"};
-  memgraph::utils::OnScopeExit replica1_directory_cleaner(
-      [&]() { std::filesystem::remove_all(replica1_storage_directory); });
-  memgraph::utils::OnScopeExit replica2_directory_cleaner(
-      [&]() { std::filesystem::remove_all(replica2_storage_directory); });
-
-  UpdatePaths(replica1_config, replica1_storage_directory);
-  UpdatePaths(replica2_config, replica2_storage_directory);
-
-  std::optional<MinMemgraph> main(main_config);
-  MinMemgraph replica1(replica1_config);
-
-  replica1.repl_handler.SetReplicationRoleReplica(ReplicationServerConfig{
-      .ip_address = local_host,
-      .port = ports[0],
-  });
-
-  MinMemgraph replica2(replica2_config);
-  replica2.repl_handler.SetReplicationRoleReplica(ReplicationServerConfig{
-      .ip_address = local_host,
-      .port = ports[1],
-  });
-
-  auto res = main->repl_handler.RegisterReplica(ReplicationClientConfig{
-      .name = replicas[0],
-      .mode = ReplicationMode::SYNC,
-      .ip_address = local_host,
-      .port = ports[0],
-  });
-  ASSERT_FALSE(res.HasError()) << (int)res.GetError();
-  res = main->repl_handler.RegisterReplica(ReplicationClientConfig{
-      .name = replicas[1],
-      .mode = ReplicationMode::SYNC,
-      .ip_address = local_host,
-      .port = ports[1],
-  });
-  ASSERT_FALSE(res.HasError()) << (int)res.GetError();
-
-  auto replica_infos = main->db.storage()->ReplicasInfo();
-
-  ASSERT_EQ(replica_infos.size(), 2);
-  ASSERT_EQ(replica_infos[0].name, replicas[0]);
-  ASSERT_EQ(replica_infos[0].endpoint.address, local_host);
-  ASSERT_EQ(replica_infos[0].endpoint.port, ports[0]);
-  ASSERT_EQ(replica_infos[1].name, replicas[1]);
-  ASSERT_EQ(replica_infos[1].endpoint.address, local_host);
-  ASSERT_EQ(replica_infos[1].endpoint.port, ports[1]);
-
-  main.reset();
-
-  MinMemgraph other_main(main_config);
-
-  replica_infos = other_main.db.storage()->ReplicasInfo();
-  ASSERT_EQ(replica_infos.size(), 2);
-  ASSERT_EQ(replica_infos[0].name, replicas[0]);
-  ASSERT_EQ(replica_infos[0].endpoint.address, local_host);
-  ASSERT_EQ(replica_infos[0].endpoint.port, ports[0]);
-  ASSERT_EQ(replica_infos[1].name, replicas[1]);
-  ASSERT_EQ(replica_infos[1].endpoint.address, local_host);
-  ASSERT_EQ(replica_infos[1].endpoint.port, ports[1]);
-}
-
 TEST_F(ReplicationTest, RestoringReplicationAtStartup) {
   auto main_config = main_conf;
   main_config.durability.restore_replication_state_on_startup = true;
 
   std::optional<MinMemgraph> main(main_config);
+  main->repl_handler.SetReplicationRoleMain(ReplicationServerConfig{
+      .ip_address = local_host,
+      .port = 10003,
+  });
+
   MinMemgraph replica1(repl_conf);
 
   replica1.repl_handler.SetReplicationRoleReplica(ReplicationServerConfig{
@@ -1062,6 +1001,80 @@ TEST_F(ReplicationTest, RestoringReplicationAtStartup) {
   ASSERT_EQ(replica_infos[0].endpoint.port, ports[1]);
 }
 
+TEST_F(ReplicationTest, RestoringReplicationAtStartupAfterDroppingReplica) {
+  auto main_config = main_conf;
+  auto replica1_config = main_conf;
+  auto replica2_config = main_conf;
+  main_config.durability.restore_replication_state_on_startup = true;
+
+  std::filesystem::path replica1_storage_directory{std::filesystem::temp_directory_path() / "replica1"};
+  std::filesystem::path replica2_storage_directory{std::filesystem::temp_directory_path() / "replica2"};
+  memgraph::utils::OnScopeExit replica1_directory_cleaner(
+      [&]() { std::filesystem::remove_all(replica1_storage_directory); });
+  memgraph::utils::OnScopeExit replica2_directory_cleaner(
+      [&]() { std::filesystem::remove_all(replica2_storage_directory); });
+
+  UpdatePaths(replica1_config, replica1_storage_directory);
+  UpdatePaths(replica2_config, replica2_storage_directory);
+
+  std::optional<MinMemgraph> main(main_config);
+  main->repl_handler.SetReplicationRoleMain(ReplicationServerConfig{
+      .ip_address = local_host,
+      .port = 10003,
+  });
+
+  MinMemgraph replica1(replica1_config);
+
+  replica1.repl_handler.SetReplicationRoleReplica(ReplicationServerConfig{
+      .ip_address = local_host,
+      .port = ports[0],
+  });
+
+  MinMemgraph replica2(replica2_config);
+  replica2.repl_handler.SetReplicationRoleReplica(ReplicationServerConfig{
+      .ip_address = local_host,
+      .port = ports[1],
+  });
+
+  auto res = main->repl_handler.RegisterReplica(ReplicationClientConfig{
+      .name = replicas[0],
+      .mode = ReplicationMode::SYNC,
+      .ip_address = local_host,
+      .port = ports[0],
+  });
+  ASSERT_FALSE(res.HasError()) << (int)res.GetError();
+  res = main->repl_handler.RegisterReplica(ReplicationClientConfig{
+      .name = replicas[1],
+      .mode = ReplicationMode::SYNC,
+      .ip_address = local_host,
+      .port = ports[1],
+  });
+  ASSERT_FALSE(res.HasError()) << (int)res.GetError();
+
+  auto replica_infos = main->db.storage()->ReplicasInfo();
+
+  ASSERT_EQ(replica_infos.size(), 2);
+  ASSERT_EQ(replica_infos[0].name, replicas[0]);
+  ASSERT_EQ(replica_infos[0].endpoint.address, local_host);
+  ASSERT_EQ(replica_infos[0].endpoint.port, ports[0]);
+  ASSERT_EQ(replica_infos[1].name, replicas[1]);
+  ASSERT_EQ(replica_infos[1].endpoint.address, local_host);
+  ASSERT_EQ(replica_infos[1].endpoint.port, ports[1]);
+
+  main.reset();
+
+  MinMemgraph other_main(main_config);
+
+  replica_infos = other_main.db.storage()->ReplicasInfo();
+  ASSERT_EQ(replica_infos.size(), 2);
+  ASSERT_EQ(replica_infos[0].name, replicas[0]);
+  ASSERT_EQ(replica_infos[0].endpoint.address, local_host);
+  ASSERT_EQ(replica_infos[0].endpoint.port, ports[0]);
+  ASSERT_EQ(replica_infos[1].name, replicas[1]);
+  ASSERT_EQ(replica_infos[1].endpoint.address, local_host);
+  ASSERT_EQ(replica_infos[1].endpoint.port, ports[1]);
+}
+
 TEST_F(ReplicationTest, AddingInvalidReplica) {
   MinMemgraph main(main_conf);