fix: Add storage connection in storage interface

src/spider/CMakeLists.txt

@@ -24,6 +24,7 @@ set(SPIDER_CORE_HEADERS
     utils/LruCache.hpp
     storage/MetadataStorage.hpp
     storage/DataStorage.hpp
+    storage/StorageConnection.hpp
     storage/MySqlConnection.hpp
     storage/MySqlStorage.hpp
     worker/FunctionManager.hpp

src/spider/client/Data.hpp

@@ -6,6 +6,7 @@
 #include <memory>
 #include <string>
 #include <utility>
+#include <variant>
 #include <vector>
 
 #include <boost/uuid/uuid.hpp>
@@ -14,6 +15,7 @@
 #include "../io/MsgPack.hpp"  // IWYU pragma: keep
 #include "../io/Serializer.hpp"
 #include "../storage/DataStorage.hpp"
+#include "../storage/MySqlConnection.hpp"
 #include "Exception.hpp"
 
 namespace spider {
@@ -64,7 +66,13 @@ class Data {
     void set_locality(std::vector<std::string> const& nodes, bool hard) {
         m_impl->set_locality(nodes);
         m_impl->set_hard_locality(hard);
-        m_data_store->set_data_locality(*m_impl);
+        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);
+        }
+        auto& conn = std::get<core::MySqlConnection>(conn_result);
+        m_data_store->set_data_locality(conn, *m_impl);
     }
 
     class Builder {
@@ -108,16 +116,22 @@ 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);
+            }
+            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(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(m_source_id, *data);
+                    err = m_data_store->add_task_data(conn, m_source_id, *data);
                     if (!err.success()) {
                         throw ConnectionException(err.description);
                     }

src/spider/client/Driver.cpp

@@ -6,6 +6,7 @@
 #include <stop_token>
 #include <string>
 #include <thread>
+#include <variant>
 
 #include <boost/uuid/random_generator.hpp>
 #include <boost/uuid/uuid.hpp>
@@ -14,6 +15,7 @@
 #include "../core/Error.hpp"
 #include "../core/KeyValueData.hpp"
 #include "../io/BoostAsio.hpp"  // IWYU pragma: keep
+#include "../storage/MySqlConnection.hpp"
 #include "../storage/MySqlStorage.hpp"
 #include "Exception.hpp"
 
@@ -26,7 +28,14 @@ Driver::Driver(std::string const& storage_url) {
     m_metadata_storage = std::make_shared<core::MySqlMetadataStorage>(storage_url);
     m_data_storage = std::make_shared<core::MySqlDataStorage>(storage_url);
 
-    core::StorageErr const err = m_metadata_storage->add_driver(core::Driver{m_id});
+    std::variant<core::MySqlConnection, core::StorageErr> conn_result
+            = core::MySqlConnection::create(storage_url);
+    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);
+
+    core::StorageErr const err = m_metadata_storage->add_driver(conn, core::Driver{m_id});
     if (!err.success()) {
         if (core::StorageErrType::DuplicateKeyErr == err.type) {
             throw DriverIdInUseException(m_id);
@@ -39,7 +48,14 @@ 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));
-            core::StorageErr const err = m_metadata_storage->update_heartbeat(m_id);
+            std::variant<core::MySqlConnection, core::StorageErr> conn_result
+                    = core::MySqlConnection::create(m_metadata_storage->get_url());
+            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);
+
+            core::StorageErr const err = m_metadata_storage->update_heartbeat(conn, m_id);
             if (!err.success()) {
                 throw ConnectionException(err.description);
             }
@@ -50,8 +66,14 @@ Driver::Driver(std::string const& storage_url) {
 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);
+    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);
 
-    core::StorageErr const err = m_metadata_storage->add_driver(core::Driver{m_id});
+    core::StorageErr const err = m_metadata_storage->add_driver(conn, core::Driver{m_id});
     if (!err.success()) {
         if (core::StorageErrType::DuplicateKeyErr == err.type) {
             throw DriverIdInUseException(m_id);
@@ -64,7 +86,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));
-            core::StorageErr const err = m_metadata_storage->update_heartbeat(m_id);
+            std::variant<core::MySqlConnection, core::StorageErr> conn_result
+                    = core::MySqlConnection::create(m_metadata_storage->get_url());
+            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);
+
+            core::StorageErr const err = m_metadata_storage->update_heartbeat(conn, m_id);
             if (!err.success()) {
                 throw ConnectionException(err.description);
             }
@@ -74,15 +103,30 @@ Driver::Driver(std::string const& storage_url, boost::uuids::uuid const id) : m_
 
 auto Driver::kv_store_insert(std::string const& key, std::string const& value) -> void {
     core::KeyValueData const kv_data{key, value, m_id};
-    core::StorageErr const err = m_data_storage->add_client_kv_data(kv_data);
+
+    std::variant<core::MySqlConnection, core::StorageErr> conn_result
+            = core::MySqlConnection::create(m_data_storage->get_url());
+    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);
+
+    core::StorageErr const err = m_data_storage->add_client_kv_data(conn, kv_data);
     if (!err.success()) {
         throw ConnectionException(err.description);
     }
 }
 
 auto Driver::kv_store_get(std::string const& key) -> std::optional<std::string> {
+    std::variant<core::MySqlConnection, core::StorageErr> conn_result
+            = core::MySqlConnection::create(m_data_storage->get_url());
+    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);
+
     std::string value;
-    core::StorageErr const err = m_data_storage->get_client_kv_data(m_id, key, &value);
+    core::StorageErr const err = m_data_storage->get_client_kv_data(conn, m_id, key, &value);
     if (!err.success()) {
         if (core::StorageErrType::KeyNotFoundErr == err.type) {
             return std::nullopt;

src/spider/client/Driver.hpp

@@ -8,6 +8,7 @@
 #include <thread>
 #include <tuple>
 #include <type_traits>
+#include <variant>
 #include <vector>
 
 #include <boost/uuid/random_generator.hpp>
@@ -17,6 +18,7 @@
 #include "../core/Error.hpp"
 #include "../core/TaskGraphImpl.hpp"
 #include "../io/Serializer.hpp"
+#include "../storage/MySqlConnection.hpp"
 #include "../worker/FunctionManager.hpp"
 #include "../worker/FunctionNameManager.hpp"
 #include "Data.hpp"
@@ -173,7 +175,13 @@ class Driver {
         graph.add_task(new_task);
         graph.add_input_task(new_task.get_id());
         graph.add_output_task(new_task.get_id());
-        core::StorageErr err = m_metadata_storage->add_job(job_id, m_id, graph);
+        std::variant<core::MySqlConnection, core::StorageErr> conn_result
+                = core::MySqlConnection::create(m_metadata_storage->get_url());
+        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);
+        core::StorageErr err = m_metadata_storage->add_job(conn, job_id, m_id, graph);
         if (!err.success()) {
             throw ConnectionException(fmt::format("Failed to start job: {}", err.description));
         }
@@ -216,8 +224,14 @@ class Driver {
         graph.m_impl->reset_ids();
         boost::uuids::random_generator gen;
         boost::uuids::uuid const job_id = gen();
+        std::variant<core::MySqlConnection, core::StorageErr> conn_result
+                = core::MySqlConnection::create(m_metadata_storage->get_url());
+        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);
         core::StorageErr const err
-                = m_metadata_storage->add_job(job_id, m_id, graph.m_impl->get_graph());
+                = m_metadata_storage->add_job(conn, job_id, m_id, graph.m_impl->get_graph());
         if (!err.success()) {
             throw ConnectionException(fmt::format("Failed to start job: {}", err.description));
         }
@@ -235,7 +249,14 @@ class Driver {
      */
     auto get_jobs() -> std::vector<boost::uuids::uuid> {
         std::vector<boost::uuids::uuid> job_ids;
-        core::StorageErr const err = m_metadata_storage->get_jobs_by_client_id(m_id, &job_ids);
+        std::variant<core::MySqlConnection, core::StorageErr> conn_result
+                = core::MySqlConnection::create(m_metadata_storage->get_url());
+        if (std::holds_alternative<spider::core::StorageErr>(conn_result)) {
+            throw ConnectionException(std::get<spider::core::StorageErr>(conn_result).description);
+        }
+        auto& conn = std::get<core::MySqlConnection>(conn_result);
+        core::StorageErr const err
+                = m_metadata_storage->get_jobs_by_client_id(conn, m_id, &job_ids);
         if (!err.success()) {
             throw ConnectionException("Failed to get jobs.");
         }

src/spider/client/Job.hpp

@@ -10,6 +10,7 @@
 #include <thread>
 #include <tuple>
 #include <utility>
+#include <variant>
 #include <vector>
 
 #include <boost/uuid/uuid.hpp>
@@ -20,7 +21,9 @@
 #include "../core/JobMetadata.hpp"
 #include "../io/MsgPack.hpp"  // IWYU pragma: keep
 #include "../storage/MetadataStorage.hpp"
+#include "../storage/MySqlConnection.hpp"
 #include "Data.hpp"
+#include "Exception.hpp"
 #include "task.hpp"
 #include "type_utils.hpp"
 
@@ -60,8 +63,15 @@ class Job {
      * @throw spider::ConnectionException
      */
     auto wait_complete() -> void {
+        std::variant<core::MySqlConnection, core::StorageErr> conn_result
+                = core::MySqlConnection::create(m_data_storage->get_url());
+        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);
+
         bool complete = false;
-        core::StorageErr err = m_metadata_storage->get_job_complete(m_id, &complete);
+        core::StorageErr err = m_metadata_storage->get_job_complete(conn, m_id, &complete);
         if (!err.success()) {
             throw ConnectionException{
                     fmt::format("Failed to get job completion status: {}", err.description)
@@ -70,7 +80,7 @@ class Job {
         while (!complete) {
             constexpr int cSleepMs = 10;
             std::this_thread::sleep_for(std::chrono::milliseconds(cSleepMs));
-            err = m_metadata_storage->get_job_complete(m_id, &complete);
+            err = m_metadata_storage->get_job_complete(conn, m_id, &complete);
             if (!err.success()) {
                 throw ConnectionException{
                         fmt::format("Failed to get job completion status: {}", err.description)
@@ -91,8 +101,15 @@ class Job {
      * @throw spider::ConnectionException
      */
     auto get_status() -> JobStatus {
+        std::variant<core::MySqlConnection, core::StorageErr> conn_result
+                = core::MySqlConnection::create(m_data_storage->get_url());
+        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);
+
         core::JobStatus status = core::JobStatus::Running;
-        core::StorageErr const err = m_metadata_storage->get_job_status(m_id, &status);
+        core::StorageErr const err = m_metadata_storage->get_job_status(conn, m_id, &status);
         if (!err.success()) {
             throw ConnectionException{fmt::format("Failed to get job status: {}", err.description)};
         }
@@ -120,8 +137,16 @@ class Job {
      * @throw spider::ConnectionException
      */
     auto get_result() -> ReturnType {
+        std::variant<core::MySqlConnection, core::StorageErr> conn_result
+                = core::MySqlConnection::create(m_data_storage->get_url());
+        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);
+
         std::vector<boost::uuids::uuid> output_task_ids;
-        core::StorageErr err = m_metadata_storage->get_job_output_tasks(m_id, &output_task_ids);
+        core::StorageErr err
+                = m_metadata_storage->get_job_output_tasks(conn, m_id, &output_task_ids);
         if (!err.success()) {
             throw ConnectionException{
                     fmt::format("Failed to get job output tasks: {}", err.description)
@@ -130,7 +155,7 @@ class Job {
         std::vector<core::Task> tasks;
         for (auto const& id : output_task_ids) {
             core::Task task{""};
-            err = m_metadata_storage->get_task(id, &task);
+            err = m_metadata_storage->get_task(conn, id, &task);
             if (!err.success()) {
                 throw ConnectionException{fmt::format("Failed to get task: {}", err.description)};
             }
@@ -161,7 +186,7 @@ class Job {
                     if (!optional_data_id.has_value()) {
                         throw ConnectionException{fmt::format("Output data ID is missing")};
                     }
-                    err = m_data_storage->get_data(optional_data_id.value(), &data);
+                    err = m_data_storage->get_data(conn, optional_data_id.value(), &data);
                     if (!err.success()) {
                         throw ConnectionException{
                                 fmt::format("Failed to get data: {}", err.description)
@@ -202,7 +227,7 @@ class Job {
                 throw ConnectionException{fmt::format("Expected one output task for job result")};
             }
             core::Task task{""};
-            err = m_metadata_storage->get_task(output_task_ids[0], &task);
+            err = m_metadata_storage->get_task(conn, output_task_ids[0], &task);
             if (!err.success()) {
                 throw ConnectionException{fmt::format("Failed to get task: {}", err.description)};
             }
@@ -220,7 +245,7 @@ class Job {
                 if (!optional_data_id.has_value()) {
                     throw ConnectionException{fmt::format("Output data ID is missing")};
                 }
-                err = m_data_storage->get_data(optional_data_id.value(), &data);
+                err = m_data_storage->get_data(conn, optional_data_id.value(), &data);
                 if (!err.success()) {
                     throw ConnectionException{fmt::format("Failed to get data: {}", err.description)
                     };