diff --git a/src/spider/client/Data.hpp b/src/spider/client/Data.hpp index b329e439..aae33e0f 100644 --- a/src/spider/client/Data.hpp +++ b/src/spider/client/Data.hpp @@ -78,6 +78,26 @@ class Data { m_data_store->set_data_locality(*conn, *m_impl); } + /** + * Sets the data as persisted, indicating the data should not be cleaned up. + * + * @throw spider::ConnectionException + */ + void set_persisted() { + m_impl->set_persisted(true); + if (nullptr != m_connection) { + m_data_store->set_data_persisted(*m_connection, *m_impl); + return; + } + std::variant, core::StorageErr> conn_result + = m_storage_factory->provide_storage_connection(); + if (std::holds_alternative(conn_result)) { + throw ConnectionException(std::get(conn_result).description); + } + auto conn = std::move(std::get>(conn_result)); + m_data_store->set_data_persisted(*conn, *m_impl); + } + class Builder { public: /** @@ -106,6 +126,16 @@ class Data { return *this; } + /** + * Sets the data as persisted, indicating the data should not be cleaned up. + * + * @return self + */ + auto set_persisted() -> Builder& { + m_persisted = true; + return *this; + } + /** * Builds the data object. * @@ -119,6 +149,7 @@ class Data { auto data = std::make_unique(std::string{buffer.data(), buffer.size()}); data->set_locality(m_nodes); data->set_hard_locality(m_hard_locality); + data->set_persisted(m_persisted); std::shared_ptr conn = m_connection; if (nullptr == conn) { std::variant, core::StorageErr> conn_result @@ -166,6 +197,7 @@ class Data { std::vector m_nodes; bool m_hard_locality = false; std::function m_cleanup_func; + bool m_persisted = false; std::shared_ptr m_data_store; std::shared_ptr m_storage_factory; diff --git a/src/spider/core/Data.hpp b/src/spider/core/Data.hpp index b5a1d556..1903aa95 100644 --- a/src/spider/core/Data.hpp +++ b/src/spider/core/Data.hpp @@ -31,11 +31,16 @@ class Data { void set_hard_locality(bool const hard) { m_hard_locality = hard; } + void set_persisted(bool const persisted) { this->m_persisted = persisted; } + + [[nodiscard]] auto is_persisted() const -> bool { return m_persisted; } + private: boost::uuids::uuid m_id; std::string m_value; std::vector m_locality; bool m_hard_locality = false; + bool m_persisted = false; void init_id() { boost::uuids::random_generator gen; diff --git a/src/spider/storage/DataStorage.hpp b/src/spider/storage/DataStorage.hpp index e2be1984..553005be 100644 --- a/src/spider/storage/DataStorage.hpp +++ b/src/spider/storage/DataStorage.hpp @@ -65,6 +65,7 @@ class DataStorage { ) -> StorageErr = 0; virtual auto set_data_locality(StorageConnection& conn, Data const& data) -> StorageErr = 0; + virtual auto set_data_persisted(StorageConnection& conn, Data const& data) -> StorageErr = 0; virtual auto remove_data(StorageConnection& conn, boost::uuids::uuid id) -> StorageErr = 0; virtual auto add_task_reference(StorageConnection& conn, boost::uuids::uuid id, boost::uuids::uuid task_id) diff --git a/src/spider/storage/mysql/MySqlStorage.cpp b/src/spider/storage/mysql/MySqlStorage.cpp index a7c321d9..74cba3a7 100644 --- a/src/spider/storage/mysql/MySqlStorage.cpp +++ b/src/spider/storage/mysql/MySqlStorage.cpp @@ -2010,13 +2010,15 @@ auto MySqlDataStorage::add_driver_data( try { std::unique_ptr statement( static_cast(conn)->prepareStatement( - "INSERT INTO `data` (`id`, `value`, `hard_locality`) VALUES(?, ?, ?)" + "INSERT INTO `data` (`id`, `value`, `hard_locality`, `persisted`) " + "VALUES(?, ?, ?, ?)" ) ); sql::bytes id_bytes = uuid_get_bytes(data.get_id()); statement->setBytes(1, &id_bytes); statement->setString(2, data.get_value()); statement->setBoolean(3, data.is_hard_locality()); + statement->setBoolean(4, data.is_persisted()); statement->executeUpdate(); for (std::string const& addr : data.get_locality()) { @@ -2058,13 +2060,15 @@ auto MySqlDataStorage::add_task_data( try { std::unique_ptr statement( static_cast(conn)->prepareStatement( - "INSERT INTO `data` (`id`, `value`, `hard_locality`) VALUES(?, ?, ?)" + "INSERT INTO `data` (`id`, `value`, `hard_locality`, `persisted`) " + "VALUES(?, ?, ?, ?)" ) ); sql::bytes id_bytes = uuid_get_bytes(data.get_id()); statement->setBytes(1, &id_bytes); statement->setString(2, data.get_value()); statement->setBoolean(3, data.is_hard_locality()); + statement->setBoolean(4, data.is_persisted()); statement->executeUpdate(); for (std::string const& addr : data.get_locality()) { @@ -2104,7 +2108,7 @@ auto MySqlDataStorage::get_data_with_locality( ) -> StorageErr { std::unique_ptr statement( static_cast(conn)->prepareStatement( - "SELECT `id`, `value`, `hard_locality` FROM `data` WHERE `id` = ?" + "SELECT `id`, `value`, `hard_locality`, `persisted` FROM `data` WHERE `id` = ?" ) ); sql::bytes id_bytes = uuid_get_bytes(id); @@ -2120,6 +2124,7 @@ auto MySqlDataStorage::get_data_with_locality( res->next(); *data = Data{id, get_sql_string(res->getString(2))}; data->set_hard_locality(res->getBoolean(3)); + data->set_persisted(res->getBoolean(4)); std::unique_ptr locality_statement( static_cast(conn)->prepareStatement( @@ -2249,6 +2254,25 @@ auto MySqlDataStorage::set_data_locality(StorageConnection& conn, Data const& da return StorageErr{}; } +auto MySqlDataStorage::set_data_persisted(StorageConnection& conn, Data const& data) -> StorageErr { + try { + sql::bytes id_bytes = uuid_get_bytes(data.get_id()); + std::unique_ptr statement( + static_cast(conn)->prepareStatement( + "UPDATE `data` SET `persisted` = ? WHERE `id` = ?" + ) + ); + statement->setBoolean(1, data.is_persisted()); + statement->setBytes(2, &id_bytes); + statement->executeUpdate(); + } catch (sql::SQLException& e) { + static_cast(conn)->rollback(); + return StorageErr{StorageErrType::OtherErr, e.what()}; + } + static_cast(conn)->commit(); + return StorageErr{}; +} + auto MySqlDataStorage::remove_data(StorageConnection& conn, boost::uuids::uuid id) -> StorageErr { try { std::unique_ptr statement( diff --git a/src/spider/storage/mysql/MySqlStorage.hpp b/src/spider/storage/mysql/MySqlStorage.hpp index 90808dac..5b1c67f0 100644 --- a/src/spider/storage/mysql/MySqlStorage.hpp +++ b/src/spider/storage/mysql/MySqlStorage.hpp @@ -163,6 +163,7 @@ class MySqlDataStorage : public DataStorage { Data* data ) -> StorageErr override; auto set_data_locality(StorageConnection& conn, Data const& data) -> StorageErr override; + auto set_data_persisted(StorageConnection& conn, Data const& data) -> StorageErr override; auto remove_data(StorageConnection& conn, boost::uuids::uuid id) -> StorageErr override; auto add_task_reference(StorageConnection& conn, boost::uuids::uuid id, boost::uuids::uuid task_id) diff --git a/tests/storage/test-DataStorage.cpp b/tests/storage/test-DataStorage.cpp index e415ed68..a6b1ee38 100644 --- a/tests/storage/test-DataStorage.cpp +++ b/tests/storage/test-DataStorage.cpp @@ -40,7 +40,7 @@ TEMPLATE_LIST_TEST_CASE( auto conn = std::move(std::get>(conn_result)); // Add driver and data - spider::core::Data const data{"value"}; + spider::core::Data data{"value"}; boost::uuids::random_generator gen; boost::uuids::uuid const driver_id = gen(); REQUIRE(metadata_storage->add_driver(*conn, spider::core::Driver{driver_id}).success()); @@ -56,6 +56,13 @@ TEMPLATE_LIST_TEST_CASE( REQUIRE(data_storage->get_data(*conn, data.get_id(), &result).success()); REQUIRE(spider::test::data_equal(data, result)); + // Set data persisted should succeed + data.set_persisted(true); + REQUIRE(data_storage->set_data_persisted(*conn, data).success()); + // Get data should match + REQUIRE(data_storage->get_data(*conn, data.get_id(), &result).success()); + REQUIRE(spider::test::data_equal(data, result)); + // Remove data should succeed REQUIRE(data_storage->remove_data(*conn, data.get_id()).success()); diff --git a/tests/utils/CoreDataUtils.hpp b/tests/utils/CoreDataUtils.hpp index 3136e6df..50d848dc 100644 --- a/tests/utils/CoreDataUtils.hpp +++ b/tests/utils/CoreDataUtils.hpp @@ -20,6 +20,10 @@ inline auto data_equal(core::Data const& d1, core::Data const& d2) -> bool { return false; } + if (d1.is_persisted() != d2.is_persisted()) { + return false; + } + return true; } } // namespace spider::test