From e1c1db015981518a08e71de69c69e40f2f7d1182 Mon Sep 17 00:00:00 2001
From: Shuhao Wu <shuhao@shuhaowu.com>
Date: Wed, 13 May 2015 16:08:16 -0400
Subject: [PATCH] Allow custom slave connection

This commit allows for custom connection to be passed when
establishing connection to slaves.

While we can assume that with traditional rails apps, a straightforward
host substitution will be suffice, this assumption cannot be made for
projects of a certain scale, where slave boxes may have custom
configurations (maybe different password/username/etc).

This commit's approach relies on the consumer of this package to specify
a lambda that takes in the slave hostname and returns a connection
object.
---
 lib/lhm/invoker.rb                    |  3 ++-
 lib/lhm/throttler/slave_lag.rb        | 14 +++++++++-----
 spec/unit/throttler/slave_lag_spec.rb | 13 +++++++++++++
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/lib/lhm/invoker.rb b/lib/lhm/invoker.rb
index 72baac1f..16a00a19 100644
--- a/lib/lhm/invoker.rb
+++ b/lib/lhm/invoker.rb
@@ -68,7 +68,8 @@ def normalize_options(options)
       end
 
       if options[:throttler]
-        options[:throttler] = Throttler::Factory.create_throttler(options[:throttler])
+        throttler_options = options[:throttler_options] || {}
+        options[:throttler] = Throttler::Factory.create_throttler(options[:throttler], throttler_options)
       else
         options[:throttler] = Lhm.throttler
       end
diff --git a/lib/lhm/throttler/slave_lag.rb b/lib/lhm/throttler/slave_lag.rb
index 425c7d97..fad48087 100644
--- a/lib/lhm/throttler/slave_lag.rb
+++ b/lib/lhm/throttler/slave_lag.rb
@@ -15,7 +15,7 @@ def initialize(options = {})
         @timeout_seconds = INITIAL_TIMEOUT
         @stride = options[:stride] || DEFAULT_STRIDE
         @allowed_lag = options[:allowed_lag] || DEFAULT_MAX_ALLOWED_LAG
-        @slave_connections = {}
+        @slave_connection = options[:slave_connection]
       end
 
       def execute
@@ -74,10 +74,14 @@ def slave_lag(slave)
       end
 
       def slave_connection(slave)
-        adapter_method = defined?(Mysql2) ? 'mysql2_connection' : 'mysql_connection'
-        config = ActiveRecord::Base.connection_pool.spec.config.dup
-        config[:host] = slave
-        ActiveRecord::Base.send(adapter_method, config)
+        unless @slave_connection
+          adapter_method = defined?(Mysql2) ? 'mysql2_connection' : 'mysql_connection'
+          config = ActiveRecord::Base.connection_pool.spec.config.dup
+          config[:host] = slave
+          ActiveRecord::Base.send(adapter_method, config)
+        else
+          @slave_connection.call(slave)
+        end
       end
 
       # This method fetch the Seconds_Behind_Master, when exec_query is no available, on AR 2.3.
diff --git a/spec/unit/throttler/slave_lag_spec.rb b/spec/unit/throttler/slave_lag_spec.rb
index dbb41f28..97d58984 100644
--- a/spec/unit/throttler/slave_lag_spec.rb
+++ b/spec/unit/throttler/slave_lag_spec.rb
@@ -62,6 +62,19 @@ def @throttler.max_current_slave_lag
     end
   end
 
+  describe '#custom_slave_connection' do
+    describe 'with a custom slave connection' do
+      before do
+        @connection = Object.new # just an object to see if the method returns what we return
+        @throttler = Lhm::Throttler::SlaveLag.new(slave_connection: lambda {|host| @connection})
+      end
+
+      it 'should use the custom slave connection' do
+        assert_equal(@connection, @throttler.send(:slave_connection, "slave.db.local"))
+      end
+    end
+  end
+
   describe '#slave_hosts' do
     describe 'with no slaves' do
       before do