diff --git a/redis/include/userver/storages/redis/transaction.hpp b/redis/include/userver/storages/redis/transaction.hpp index b881712c4570..126968b858a5 100644 --- a/redis/include/userver/storages/redis/transaction.hpp +++ b/redis/include/userver/storages/redis/transaction.hpp @@ -9,6 +9,7 @@ #include #include #include +#include USERVER_NAMESPACE_BEGIN @@ -413,7 +414,21 @@ class Transaction { virtual RequestJsonMset JsonMset(std::vector key_path_values) = 0; + /// @brief Execute a custom Redis command. + /// @param key_index Index of the key in the args vector used to determine the shard + template + RequestGeneric GenericCommand(std::string command, std::vector args, size_t key_index) { + return RequestGeneric{GenericCommon(std::move(command), std::move(args), key_index)}; + } + // end of redis commands + +protected: + virtual RequestGenericCommon GenericCommon( + std::string command, + std::vector args, + size_t key_index + ) = 0; }; using TransactionPtr = std::unique_ptr; diff --git a/redis/src/storages/redis/transaction_impl.cpp b/redis/src/storages/redis/transaction_impl.cpp index d85a86a5d03a..b6b7f2b487d0 100644 --- a/redis/src/storages/redis/transaction_impl.cpp +++ b/redis/src/storages/redis/transaction_impl.cpp @@ -1041,6 +1041,15 @@ RequestJsonMset TransactionImpl::JsonMset(std::vector key_path return AddCmd("json.mset", true, std::move(args)); } +RequestGenericCommon TransactionImpl::GenericCommon( + std::string command, + std::vector args, + size_t key_index +) { + UpdateShard(args.at(key_index)); + return AddCmd(std::move(command), true, std::move(args)); +} + // end of redis commands void TransactionImpl::UpdateShard(const std::string& key) { diff --git a/redis/src/storages/redis/transaction_impl.hpp b/redis/src/storages/redis/transaction_impl.hpp index 8e1e8f2530dc..7bbfe11fa1a8 100644 --- a/redis/src/storages/redis/transaction_impl.hpp +++ b/redis/src/storages/redis/transaction_impl.hpp @@ -417,6 +417,8 @@ class TransactionImpl final : public Transaction { RequestJsonMset JsonMset(std::vector key_path_values) override; + RequestGenericCommon GenericCommon(std::string command, std::vector args, size_t key_index) override; + // end of redis commands private: diff --git a/redis/testing/include/userver/storages/redis/mock_transaction.hpp b/redis/testing/include/userver/storages/redis/mock_transaction.hpp index 6174ef123f3e..ca002d4b7abd 100644 --- a/redis/testing/include/userver/storages/redis/mock_transaction.hpp +++ b/redis/testing/include/userver/storages/redis/mock_transaction.hpp @@ -363,6 +363,8 @@ class MockTransaction final : public Transaction { RequestJsonMset JsonMset(std::vector key_path_values) override; + RequestGenericCommon GenericCommon(std::string command, std::vector args, size_t key_index) override; + // end of redis commands private: diff --git a/redis/testing/include/userver/storages/redis/mock_transaction_impl_base.hpp b/redis/testing/include/userver/storages/redis/mock_transaction_impl_base.hpp index ccfca6788fa0..88accc5fa7e7 100644 --- a/redis/testing/include/userver/storages/redis/mock_transaction_impl_base.hpp +++ b/redis/testing/include/userver/storages/redis/mock_transaction_impl_base.hpp @@ -363,6 +363,8 @@ class MockTransactionImplBase { virtual RequestJsonMset JsonMset(std::vector key_path_values); + virtual RequestGenericCommon GenericCommon(std::string command, std::vector args, size_t key_index); + // end of redis commands }; diff --git a/redis/testing/src/storages/redis/mock_transaction.cpp b/redis/testing/src/storages/redis/mock_transaction.cpp index aa77a1de7d5d..25140c81f4b4 100644 --- a/redis/testing/src/storages/redis/mock_transaction.cpp +++ b/redis/testing/src/storages/redis/mock_transaction.cpp @@ -873,6 +873,15 @@ RequestJsonMset MockTransaction::JsonMset(std::vector key_path return AddSubrequest(impl_->JsonMset(std::move(key_path_values))); } +RequestGenericCommon MockTransaction::GenericCommon( + std::string command, + std::vector args, + size_t key_index +) { + UpdateShard(args.at(key_index)); + return AddSubrequest(impl_->GenericCommon(std::move(command), std::move(args), key_index)); +} + // end of redis commands void MockTransaction::UpdateShard(const std::string& key) { diff --git a/redis/testing/src/storages/redis/mock_transaction_impl_base.cpp b/redis/testing/src/storages/redis/mock_transaction_impl_base.cpp index 7ad4f58e8665..971114a4aa8c 100644 --- a/redis/testing/src/storages/redis/mock_transaction_impl_base.cpp +++ b/redis/testing/src/storages/redis/mock_transaction_impl_base.cpp @@ -665,6 +665,14 @@ RequestJsonMset MockTransactionImplBase::JsonMset(std::vector AbortWithStacktrace("Redis method not mocked"); } +RequestGenericCommon MockTransactionImplBase::GenericCommon( + std::string /*command*/, + std::vector /*args*/, + size_t /*key_index*/ +) { + AbortWithStacktrace("Redis method not mocked"); +} + // end of redis commands } // namespace storages::redis