Early check connection blacklisting

lib/active_record/connection_adapters/makara_abstract_adapter.rb

@@ -95,10 +95,22 @@ def custom_error_message?(connection, message)
       control_method :close, :steal!, :expire, :lease, :in_use?, :owner, :schema_cache, :pool=, :pool,
          :schema_cache=, :lock, :seconds_idle, :==
 
-      SQL_MASTER_MATCHERS           = [/\A\s*select.+for update\Z/i, /select.+lock in share mode\Z/i, /\A\s*select.+(nextval|currval|lastval|get_lock|release_lock|pg_advisory_lock|pg_advisory_unlock)\(/i].map(&:freeze).freeze
-      SQL_SLAVE_MATCHERS            = [/\A\s*(select|with.+\)\s*select)\s/i].map(&:freeze).freeze
-      SQL_ALL_MATCHERS              = [/\A\s*set\s/i].map(&:freeze).freeze
-      SQL_SKIP_STICKINESS_MATCHERS  = [/\A\s*show\s([\w]+\s)?(field|table|database|schema|view|index)(es|s)?/i, /\A\s*(set|describe|explain|pragma)\s/i].map(&:freeze).freeze
+      SQL_MASTER_MATCHERS           = [
+        /\A\s*select.+for update\Z/i,
+        /select.+lock in share mode\Z/i,
+        /\A\s*select.+(nextval|currval|lastval|get_lock|release_lock|pg_advisory_lock|pg_advisory_unlock|pg_try_advisory_lock)\(/i,
+        /\A\s*select.+from pg_locks/i
+      ].map(&:freeze).freeze
+      SQL_SLAVE_MATCHERS            = [
+        /\A\s*(select|with.+\)\s*select)\s/i
+      ].map(&:freeze).freeze
+      SQL_ALL_MATCHERS              = [
+        /\A\s*set\s/i
+      ].map(&:freeze).freeze
+      SQL_SKIP_STICKINESS_MATCHERS  = [
+        /\A\s*show\s([\w]+\s)?(field|table|database|schema|view|index)(es|s)?/i,
+        /\A\s*(set|describe|explain|pragma)\s/i
+      ].map(&:freeze).freeze
 
       def sql_master_matchers
         SQL_MASTER_MATCHERS

lib/makara/connection_wrapper.rb

@@ -71,6 +71,10 @@ def _makara_connected?
     end
 
     def _makara_connection
+      if _makara_blacklisted? && initial_error
+        raise Makara::Errors::BlacklistConnection.new(self, initial_error)
+      end
+
       current = @connection
 
       if current
@@ -80,8 +84,9 @@ def _makara_connection
 
         # Already wrapped because of initial failure
         if new_connection.is_a?(Makara::ConnectionWrapper)
+          initial_error = new_connection.initial_error
           _makara_blacklist!
-          raise Makara::Errors::BlacklistConnection.new(self, new_connection.initial_error)
+          raise Makara::Errors::BlacklistConnection.new(self, initial_error)
         else
           @connection = new_connection
           _makara_decorate_connection(new_connection)

spec/connection_wrapper_spec.rb

@@ -39,4 +39,43 @@
       end
     end
   end
+
+  describe '#_makara_connection' do
+    it 'return connection when successfully connected' do
+      expect(subject._makara_connection).to eq(connection)
+    end
+
+    it 'raise error when blacklisted with initial_error' do
+      expect(subject).to receive(:initial_error).and_return(StandardError.new('some connection issue')).twice # master + slave
+      subject._makara_blacklist!
+
+      expect(proxy).not_to receive(:graceful_connection_for)
+
+      expect{ subject._makara_connection }.to raise_error(Makara::Errors::BlacklistConnection)
+    end
+
+    context 'not connected' do
+      it 'return connection when re-connecting successfully' do
+        fake_connection = FakeConnection.new({:master_ttl=>5, :blacklist_duration=>5, :sticky=>true, :name=>"master/1"})
+
+        subject.instance_variable_set(:@connection, nil)
+        expect(proxy).to receive(:graceful_connection_for).and_return(fake_connection)
+
+        expect(subject._makara_connection).to eq(fake_connection)
+      end
+
+      it 'raise error when unable to connect' do
+        fake_connection = FakeConnection.new({:master_ttl=>5, :blacklist_duration=>5, :sticky=>true, :name=>"master/1"})
+        connection_error = StandardError.new('some connection issue')
+        expect(fake_connection).to receive(:is_a?).with(Makara::ConnectionWrapper).and_return(true)
+        expect(fake_connection).to receive(:initial_error).and_return(connection_error)
+
+        subject.instance_variable_set(:@connection, nil)
+        expect(proxy).to receive(:graceful_connection_for).and_return(fake_connection)
+
+        expect(subject).to receive(:_makara_blacklist!)
+        expect{ subject._makara_connection }.to raise_error(Makara::Errors::BlacklistConnection)
+      end
+    end
+  end
 end