refactor: Add abstract factory pattern to storage

src/spider/CMakeLists.txt

@@ -1,6 +1,8 @@
 # set variable as CACHE INTERNAL to access it from other scope
 set(SPIDER_CORE_SOURCES
     storage/mysql/MySqlConnection.cpp
+    storage/mysql/MySqlStorageFactory.cpp
+    storage/mysql/MySqlJobSubmissionBatch.cpp
     storage/mysql/MySqlStorage.cpp
     worker/FunctionManager.cpp
     worker/FunctionNameManager.cpp
@@ -27,9 +29,11 @@ set(SPIDER_CORE_HEADERS
     storage/StorageConnection.hpp
     storage/mysql/mysql_stmt.hpp
     storage/mysql/MySqlConnection.hpp
+    storage/mysql/MySqlStorageFactory.hpp
     storage/mysql/MySqlStorage.hpp
     storage/mysql/MySqlJobSubmissionBatch.hpp
     storage/JobSubmissionBatch.hpp
+    storage/StorageFactory.hpp
     worker/FunctionManager.hpp
     worker/FunctionNameManager.hpp
     CACHE INTERNAL
@@ -91,6 +95,7 @@ target_link_libraries(
         Boost::program_options
         Boost::system
         ${CMAKE_DL_LIBS}
+        fmt::fmt
         spdlog::spdlog
 )
 
@@ -107,6 +112,7 @@ target_link_libraries(
         Boost::program_options
         Boost::system
         ${CMAKE_DL_LIBS}
+        fmt::fmt
         spdlog::spdlog
 )
 add_dependencies(spider_worker spider_task_executor)
@@ -132,6 +138,7 @@ target_link_libraries(
         Boost::headers
         Boost::program_options
         absl::flat_hash_map
+        fmt::fmt
         spdlog::spdlog
 )
 

src/spider/client/Data.hpp

@@ -15,7 +15,8 @@
 #include "../io/MsgPack.hpp"  // IWYU pragma: keep
 #include "../io/Serializer.hpp"
 #include "../storage/DataStorage.hpp"
-#include "../storage/mysql/MySqlConnection.hpp"
+#include "../storage/StorageConnection.hpp"
+#include "../storage/StorageFactory.hpp"
 #include "Exception.hpp"
 
 namespace spider {
@@ -66,13 +67,17 @@ class Data {
     void set_locality(std::vector<std::string> const& nodes, bool hard) {
         m_impl->set_locality(nodes);
         m_impl->set_hard_locality(hard);
-        std::variant<core::MySqlConnection, core::StorageErr> conn_result
-                = core::MySqlConnection::create(m_data_store->get_url());
+        if (nullptr != m_connection) {
+            m_data_store->set_data_locality(*m_connection, *m_impl);
+            return;
+        }
+        std::variant<std::unique_ptr<core::StorageConnection>, core::StorageErr> conn_result
+                = m_storage_factory->provide_storage_connection();
         if (std::holds_alternative<core::StorageErr>(conn_result)) {
             throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
         }
-        auto& conn = std::get<core::MySqlConnection>(conn_result);
-        m_data_store->set_data_locality(conn, *m_impl);
+        auto conn = std::move(std::get<std::unique_ptr<core::StorageConnection>>(conn_result));
+        m_data_store->set_data_locality(*conn, *m_impl);
     }
 
     class Builder {
@@ -116,28 +121,31 @@ class Data {
             auto data = std::make_unique<core::Data>(std::string{buffer.data(), buffer.size()});
             data->set_locality(m_nodes);
             data->set_hard_locality(m_hard_locality);
-            std::variant<core::MySqlConnection, core::StorageErr> conn_result
-                    = core::MySqlConnection::create(m_data_store->get_url());
-            if (std::holds_alternative<core::StorageErr>(conn_result)) {
-                throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
+            std::shared_ptr<core::StorageConnection> conn = m_connection;
+            if (nullptr == conn) {
+                std::variant<std::unique_ptr<core::StorageConnection>, core::StorageErr> conn_result
+                        = m_storage_factory->provide_storage_connection();
+                if (std::holds_alternative<core::StorageErr>(conn_result)) {
+                    throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
+                }
+                conn = std::move(std::get<std::unique_ptr<core::StorageConnection>>(conn_result));
             }
-            auto& conn = std::get<core::MySqlConnection>(conn_result);
             core::StorageErr err;
             switch (m_data_source) {
                 case DataSource::Driver:
-                    err = m_data_store->add_driver_data(conn, m_source_id, *data);
+                    err = m_data_store->add_driver_data(*conn, m_source_id, *data);
                     if (!err.success()) {
                         throw ConnectionException(err.description);
                     }
                     break;
                 case DataSource::TaskContext:
-                    err = m_data_store->add_task_data(conn, m_source_id, *data);
+                    err = m_data_store->add_task_data(*conn, m_source_id, *data);
                     if (!err.success()) {
                         throw ConnectionException(err.description);
                     }
                     break;
             }
-            return Data{std::move(data), m_data_store};
+            return Data{std::move(data), m_data_store, m_storage_factory, m_connection};
         }
 
     private:
@@ -148,16 +156,32 @@ class Data {
 
         Builder(std::shared_ptr<core::DataStorage> data_store,
                 boost::uuids::uuid const source_id,
-                DataSource const data_source)
+                DataSource const data_source,
+                std::shared_ptr<core::StorageFactory> storage_factory)
                 : m_data_store{std::move(data_store)},
                   m_source_id{source_id},
-                  m_data_source{data_source} {}
+                  m_data_source{data_source},
+                  m_storage_factory{std::move(storage_factory)} {}
+
+        Builder(std::shared_ptr<core::DataStorage> data_store,
+                boost::uuids::uuid const source_id,
+                DataSource const data_source,
+                std::shared_ptr<core::StorageFactory> storage_factory,
+                std::shared_ptr<core::StorageConnection> connection)
+                : m_data_store{std::move(data_store)},
+                  m_source_id{source_id},
+                  m_data_source{data_source},
+                  m_storage_factory{std::move(storage_factory)},
+                  m_connection{std::move(connection)} {}
 
         std::vector<std::string> m_nodes;
         bool m_hard_locality = false;
         std::function<void(T const&)> m_cleanup_func;
 
         std::shared_ptr<core::DataStorage> m_data_store;
+        std::shared_ptr<core::StorageFactory> m_storage_factory;
+        std::shared_ptr<core::StorageConnection> m_connection = nullptr;
+
         boost::uuids::uuid m_source_id;
         DataSource m_data_source;
 
@@ -168,14 +192,28 @@ class Data {
     Data() = default;
 
 private:
-    Data(std::unique_ptr<core::Data> impl, std::shared_ptr<core::DataStorage> data_store)
+    Data(std::unique_ptr<core::Data> impl,
+         std::shared_ptr<core::DataStorage> data_store,
+         std::shared_ptr<core::StorageFactory> storage_factory)
+            : m_impl{std::move(impl)},
+              m_data_store{std::move(data_store)},
+              m_storage_factory{std::move(storage_factory)} {}
+
+    Data(std::unique_ptr<core::Data> impl,
+         std::shared_ptr<core::DataStorage> data_store,
+         std::shared_ptr<core::StorageFactory> storage_factory,
+         std::shared_ptr<core::StorageConnection> connection)
             : m_impl{std::move(impl)},
-              m_data_store{std::move(data_store)} {}
+              m_data_store{std::move(data_store)},
+              m_storage_factory{std::move(storage_factory)},
+              m_connection{std::move(connection)} {}
 
     [[nodiscard]] auto get_impl() const -> std::unique_ptr<core::Data> const& { return m_impl; }
 
     std::unique_ptr<core::Data> m_impl;
     std::shared_ptr<core::DataStorage> m_data_store;
+    std::shared_ptr<core::StorageFactory> m_storage_factory;
+    std::shared_ptr<core::StorageConnection> m_connection = nullptr;
 
     friend class core::DataImpl;
     friend class core::TaskGraphImpl;

src/spider/client/Driver.cpp

@@ -16,27 +16,26 @@
 #include "../core/Error.hpp"
 #include "../core/KeyValueData.hpp"
 #include "../io/BoostAsio.hpp"  // IWYU pragma: keep
-#include "../storage/mysql/MySqlConnection.hpp"
-#include "../storage/mysql/MySqlStorage.hpp"
+#include "../storage/mysql/MySqlStorageFactory.hpp"
+#include "../storage/StorageConnection.hpp"
 #include "Exception.hpp"
 
 namespace spider {
 
-Driver::Driver(std::string const& storage_url) {
+Driver::Driver(std::string const& storage_url)
+        : m_storage_factory{std::make_shared<core::MySqlStorageFactory>(storage_url)} {
     boost::uuids::random_generator gen;
     m_id = gen();
 
-    m_metadata_storage = std::make_shared<core::MySqlMetadataStorage>(storage_url);
-    m_data_storage = std::make_shared<core::MySqlDataStorage>(storage_url);
+    m_metadata_storage = m_storage_factory->provide_metadata_storage();
+    m_data_storage = m_storage_factory->provide_data_storage();
 
-    std::variant<core::MySqlConnection, core::StorageErr> conn_result
-            = core::MySqlConnection::create(storage_url);
+    std::variant<std::unique_ptr<core::StorageConnection>, core::StorageErr> conn_result
+            = m_storage_factory->provide_storage_connection();
     if (std::holds_alternative<core::StorageErr>(conn_result)) {
         throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
     }
-    m_conn = std::make_shared<core::MySqlConnection>(
-            std::get<core::MySqlConnection>(std::move(conn_result))
-    );
+    m_conn = std::move(std::get<std::unique_ptr<core::StorageConnection>>(conn_result));
 
     core::StorageErr const err = m_metadata_storage->add_driver(*m_conn, core::Driver{m_id});
     if (!err.success()) {
@@ -51,32 +50,33 @@ Driver::Driver(std::string const& storage_url) {
     m_heartbeat_thread = std::jthread([this](std::stop_token stoken) {
         while (!stoken.stop_requested()) {
             std::this_thread::sleep_for(std::chrono::seconds(1));
-            std::variant<core::MySqlConnection, core::StorageErr> conn_result
-                    = core::MySqlConnection::create(m_metadata_storage->get_url());
+            std::variant<std::unique_ptr<core::StorageConnection>, core::StorageErr> conn_result
+                    = m_storage_factory->provide_storage_connection();
             if (std::holds_alternative<core::StorageErr>(conn_result)) {
                 throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
             }
-            auto& conn = std::get<core::MySqlConnection>(conn_result);
+            auto conn = std::move(std::get<std::unique_ptr<core::StorageConnection>>(conn_result));
 
-            core::StorageErr const err = m_metadata_storage->update_heartbeat(conn, m_id);
+            core::StorageErr const err = m_metadata_storage->update_heartbeat(*conn, m_id);
             if (!err.success()) {
                 throw ConnectionException(err.description);
             }
         }
     });
 }
 
-Driver::Driver(std::string const& storage_url, boost::uuids::uuid const id) : m_id{id} {
-    m_metadata_storage = std::make_shared<core::MySqlMetadataStorage>(storage_url);
-    m_data_storage = std::make_shared<core::MySqlDataStorage>(storage_url);
-    std::variant<core::MySqlConnection, core::StorageErr> conn_result
-            = core::MySqlConnection::create(storage_url);
+Driver::Driver(std::string const& storage_url, boost::uuids::uuid const id)
+        : m_id{id},
+          m_storage_factory{std::make_shared<core::MySqlStorageFactory>(storage_url)} {
+    m_metadata_storage = m_storage_factory->provide_metadata_storage();
+    m_data_storage = m_storage_factory->provide_data_storage();
+
+    std::variant<std::unique_ptr<core::StorageConnection>, core::StorageErr> conn_result
+            = m_storage_factory->provide_storage_connection();
     if (std::holds_alternative<core::StorageErr>(conn_result)) {
         throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
     }
-    m_conn = std::make_shared<core::MySqlConnection>(
-            std::get<core::MySqlConnection>(std::move(conn_result))
-    );
+    m_conn = std::move(std::get<std::unique_ptr<core::StorageConnection>>(conn_result));
 
     core::StorageErr const err = m_metadata_storage->add_driver(*m_conn, core::Driver{m_id});
     if (!err.success()) {
@@ -91,14 +91,14 @@ Driver::Driver(std::string const& storage_url, boost::uuids::uuid const id) : m_
     m_heartbeat_thread = std::jthread([this](std::stop_token stoken) {
         while (!stoken.stop_requested()) {
             std::this_thread::sleep_for(std::chrono::seconds(1));
-            std::variant<core::MySqlConnection, core::StorageErr> conn_result
-                    = core::MySqlConnection::create(m_metadata_storage->get_url());
+            std::variant<std::unique_ptr<core::StorageConnection>, core::StorageErr> conn_result
+                    = m_storage_factory->provide_storage_connection();
             if (std::holds_alternative<core::StorageErr>(conn_result)) {
                 throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
             }
-            auto& conn = std::get<core::MySqlConnection>(conn_result);
+            auto conn = std::move(std::get<std::unique_ptr<core::StorageConnection>>(conn_result));
 
-            core::StorageErr const err = m_metadata_storage->update_heartbeat(conn, m_id);
+            core::StorageErr const err = m_metadata_storage->update_heartbeat(*conn, m_id);
             if (!err.success()) {
                 throw ConnectionException(err.description);
             }

src/spider/client/Driver.hpp

@@ -18,9 +18,8 @@
 #include "../core/TaskGraphImpl.hpp"
 #include "../io/Serializer.hpp"
 #include "../storage/JobSubmissionBatch.hpp"
-#include "../storage/mysql/MySqlConnection.hpp"
-#include "../storage/mysql/MySqlJobSubmissionBatch.hpp"
 #include "../storage/StorageConnection.hpp"
+#include "../storage/StorageFactory.hpp"
 #include "../worker/FunctionManager.hpp"
 #include "../worker/FunctionNameManager.hpp"
 #include "Data.hpp"
@@ -83,7 +82,13 @@ class Driver {
     template <Serializable T>
     auto get_data_builder() -> Data<T>::Builder {
         using DataBuilder = typename Data<T>::Builder;
-        return DataBuilder{m_data_storage, m_id, DataBuilder::DataSource::Driver};
+        return DataBuilder{
+                m_data_storage,
+                m_id,
+                DataBuilder::DataSource::Driver,
+                m_storage_factory,
+                m_conn
+        };
     }
 
     /**
@@ -147,10 +152,7 @@ class Driver {
         if (nullptr != m_batch) {
             return;
         }
-        m_batch = std::make_shared<core::MySqlJobSubmissionBatch>(
-                // NOLINTNEXTLINE(cppcoreguidelines-pro-type-static-cast-downcast)
-                static_cast<core::MySqlConnection&>(*m_conn)
-        );
+        m_batch = m_storage_factory->provide_job_submission_batch(*m_conn);
     }
 
     /**
@@ -224,7 +226,13 @@ class Driver {
             }
         }
 
-        return Job<ReturnType>{job_id, m_metadata_storage, m_data_storage, m_conn};
+        return Job<ReturnType>{
+                job_id,
+                m_metadata_storage,
+                m_data_storage,
+                m_storage_factory,
+                m_conn
+        };
     }
 
     /**
@@ -268,7 +276,13 @@ class Driver {
             throw ConnectionException(fmt::format("Failed to start job: {}", err.description));
         }
 
-        return Job<ReturnType>{job_id, m_metadata_storage, m_data_storage, m_conn};
+        return Job<ReturnType>{
+                job_id,
+                m_metadata_storage,
+                m_data_storage,
+                m_storage_factory,
+                m_conn
+        };
     }
 
     /**
@@ -293,6 +307,7 @@ class Driver {
     boost::uuids::uuid m_id;
     std::shared_ptr<core::MetadataStorage> m_metadata_storage;
     std::shared_ptr<core::DataStorage> m_data_storage;
+    std::shared_ptr<core::StorageFactory> m_storage_factory;
     std::shared_ptr<core::StorageConnection> m_conn;
     std::shared_ptr<core::JobSubmissionBatch> m_batch{nullptr};
     std::jthread m_heartbeat_thread;