feat: support transaction (#277)

Author: supercaracal

README.md

@@ -113,6 +113,7 @@ The following methods are able to be used like `redis-client`.
 * `#hscan`
 * `#zscan`
 * `#pipelined`
+* `#multi`
 * `#pubsub`
 * `#close`
 

lib/redis_client/cluster.rb

@@ -4,6 +4,7 @@
 require 'redis_client/cluster/pipeline'
 require 'redis_client/cluster/pub_sub'
 require 'redis_client/cluster/router'
+require 'redis_client/cluster/transaction'
 
 class RedisClient
   class Cluster
@@ -88,6 +89,13 @@ def pipelined
       pipeline.execute
     end
 
+    def multi(watch: nil, &block)
+      ::RedisClient::Cluster::Transaction
+        .new(@router, @command_builder)
+        .find_node(&block)
+        .multi(watch: watch, &block)
+    end
+
     def pubsub
       ::RedisClient::Cluster::PubSub.new(@router, @command_builder)
     end

lib/redis_client/cluster/router.rb

@@ -179,6 +179,12 @@ def find_node_key(command, seed: nil)
         end
       end
 
+      def find_primary_node_key(command)
+        key = @command.extract_first_key(command)
+        slot = key.empty? ? nil : ::RedisClient::Cluster::KeySlotConverter.convert(key)
+        @node.find_node_key_of_primary(slot)
+      end
+
       def find_node(node_key, retry_count: 3)
         @node.find_by(node_key)
       rescue ::RedisClient::Cluster::Node::ReloadNeeded

lib/redis_client/cluster/transaction.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require 'redis_client'
+
+class RedisClient
+  class Cluster
+    class Transaction
+      ConsistencyError = Class.new(::RedisClient::Error)
+
+      def initialize(router, command_builder)
+        @router = router
+        @command_builder = command_builder
+        @node_key = nil
+      end
+
+      def call(*command, **kwargs, &_)
+        command = @command_builder.generate(command, kwargs)
+        ensure_node_key(command)
+      end
+
+      def call_v(command, &_)
+        command = @command_builder.generate(command)
+        ensure_node_key(command)
+      end
+
+      def call_once(*command, **kwargs, &_)
+        command = @command_builder.generate(command, kwargs)
+        ensure_node_key(command)
+      end
+
+      def call_once_v(command, &_)
+        command = @command_builder.generate(command)
+        ensure_node_key(command)
+      end
+
+      def find_node
+        yield self
+        raise ArgumentError, 'empty transaction' if @node_key.nil?
+
+        @router.find_node(@node_key)
+      end
+
+      private
+
+      def ensure_node_key(command)
+        node_key = @router.find_primary_node_key(command)
+        raise ConsistencyError, "Client couldn't determine the node to be executed the transaction by: #{command}" if node_key.nil?
+
+        @node_key ||= node_key
+        raise ConsistencyError, "The transaction should be done for single node: #{@node_key}, #{node_key}" if node_key != @node_key
+
+        nil
+      end
+    end
+  end
+end

test/redis_client/test_cluster.rb

@@ -187,6 +187,61 @@ def test_pubsub_without_subscription
         pubsub.close
       end
 
+      def test_transaction_with_single_key
+        want = ['OK', 1, 2, '2']
+        got = @client.multi do |t|
+          t.call('SET', 'counter', '0')
+          t.call('INCR', 'counter')
+          t.call('INCR', 'counter')
+          t.call('GET', 'counter')
+        end
+
+        assert_equal(want, got)
+      end
+
+      def test_transaction_with_multiple_key
+        assert_raises(::RedisClient::Cluster::Transaction::ConsistencyError) do
+          @client.multi do |t|
+            t.call('SET', 'key1', '1')
+            t.call('SET', 'key2', '2')
+            t.call('SET', 'key3', '3')
+          end
+        end
+
+        (1..3).each do |i|
+          assert_nil(@client.call('GET', "key#{i}"))
+        end
+      end
+
+      def test_transaction_with_empty_block
+        assert_raises(ArgumentError) { @client.multi {} }
+      end
+
+      def test_transaction_with_hashtag
+        want = ['OK', 'OK', %w[1 2 3 4]]
+        got = @client.multi do |t|
+          t.call('MSET', '{key}1', '1', '{key}2', '2')
+          t.call('MSET', '{key}3', '3', '{key}4', '4')
+          t.call('MGET', '{key}1', '{key}2', '{key}3', '{key}4')
+        end
+
+        assert_equal(want, got)
+      end
+
+      def test_transaction_without_hashtag
+        assert_raises(::RedisClient::Cluster::Transaction::ConsistencyError) do
+          @client.multi do |t|
+            t.call('MSET', 'key1', '1', 'key2', '2')
+            t.call('MSET', 'key3', '3', 'key4', '4')
+            t.call('MGET', 'key1', 'key2', 'key3', 'key4')
+          end
+        end
+
+        (1..4).each do |i|
+          assert_nil(@client.call('GET', "key#{i}"))
+        end
+      end
+
       def test_pubsub_with_wrong_command
         pubsub = @client.pubsub
         assert_nil(pubsub.call('SUBWAY'))