Add auto retry mechanism on connection failure

lib/cequel/metal/keyspace.rb

@@ -21,6 +21,8 @@ class Keyspace
       attr_reader :hosts
       # @return Integer port to connect to Cassandra nodes on
       attr_reader :port
+      # @return Integer maximum number of retries to reconnect to Cassandra
+      attr_reader :max_retries
       # @return [Symbol] the default consistency for queries in this keyspace
       # @since 1.1.0
       attr_writer :default_consistency
@@ -98,6 +100,8 @@ def initialize(configuration={})
       #   single Cassandra instance to connect to
       # @option configuration [Integer] :port (9042) port on which to connect
       #   to all specified hosts
+      # @option configuration [Integer] :max-retries maximum number of retries
+      #   on connection failure
       # @option configuration [Array<String>] :hosts list of Cassandra
       #   instances to connect to (hostnames only)
       # @option configuration [String] :username user to auth with (leave blank
@@ -115,9 +119,11 @@ def configure(configuration = {})
         @configuration = configuration
 
         @hosts, @port = extract_hosts_and_port(configuration)
-        @credentials = extract_credentials(configuration)
+        @credentials  = extract_credentials(configuration)
+        @max_retries  = extract_max_retries(configuration)
 
         @name = configuration[:keyspace]
+
         # reset the connections
         clear_active_connections!
       end
@@ -153,14 +159,27 @@ def client
       #
       # Execute a CQL query in this keyspace
       #
+      #   If a connection error occurs, will retry a maximum number of
+      #   time (default 3) before re-raising the original connection
+      #   error.
+      #
       # @param statement [String] CQL string
       # @param bind_vars [Object] values for bind variables
       # @return [Enumerable] the results of the query
       #
       # @see #execute_with_consistency
       #
       def execute(statement, *bind_vars)
-        execute_with_consistency(statement, bind_vars, default_consistency)
+        retries = max_retries
+
+        begin
+          execute_with_consistency(statement, bind_vars, default_consistency)
+        rescue Cql::NotConnectedError, Ione::Io::ConnectionError => e
+          @raw_client = nil
+          raise e if retries < 0
+          retries -= 1
+          retry
+        end
       end
 
       #
@@ -269,6 +288,10 @@ def extract_hosts_and_port(configuration)
       def extract_credentials(configuration)
         configuration.slice(:username, :password).presence
       end
+
+      def extract_max_retries(configuration)
+        configuration.fetch(:"max-retries", 3)
+      end
     end
   end
 end

spec/examples/metal/keyspace_spec.rb

@@ -82,7 +82,36 @@
                            port: Cequel::SpecSupport::Helpers.port,
                            keyspace: "totallymadeup"
 
-      expect(nonexistent_keyspace.exists?).to be_false
+      expect(nonexistent_keyspace.exists?).to be false
+    end
+  end
+
+  describe "#execute" do
+    let(:statement) { "SELECT id FROM posts" }
+
+    context "without a connection error" do
+      it "executes a CQL query" do
+        expect { cequel.execute(statement) }.not_to raise_error
+      end
+    end
+
+    context "with a connection error" do
+      after(:each) do
+        RSpec::Mocks.proxy_for(cequel).reset
+      end
+
+      it "reconnects to cassandra with a new client after first failed connection" do
+        allow(cequel).to receive(:execute_with_consistency)
+          .with(statement, [], nil)
+          .and_raise(Ione::Io::ConnectionError)
+          .once
+
+        expect(cequel).to receive(:execute_with_consistency)
+          .with(statement, [], :quorum)
+          .once
+
+        expect { cequel.execute(statement) }.not_to raise_error
+      end
     end
   end
 end

templates/config/cequel.yml

@@ -3,11 +3,13 @@ development:
   host: '127.0.0.1'
   port: 9042
   keyspace: <%= app_name %>_development
+  max-retries: 3
 
 test:
   host: '127.0.0.1'
   port: 9042
   keyspace: <%= app_name %>_test
+  max-retries: 3
 
 production:
   hosts:
@@ -18,3 +20,4 @@ production:
   keyspace: <%= app_name %>_production
   username: 'myappuser'
   password: 'password1'
+  max-retries: 3