lhm-shopify 3.4.2 → 3.5.3

data/spec/integration/sql_retry/retry_with_proxysql_spec.rb

@@ -0,0 +1,108 @@
+require 'minitest/autorun'
+require 'mysql2'
+require 'lhm'
+require 'toxiproxy'
+
+require 'integration/sql_retry/lock_wait_timeout_test_helper'
+require 'integration/sql_retry/db_connection_helper'
+require 'integration/sql_retry/proxysql_helper'
+require 'integration/toxiproxy_helper'
+
+describe Lhm::SqlRetry, "ProxiSQL tests for LHM retry" do
+  include ToxiproxyHelper
+
+  before(:each) do
+    @old_logger = Lhm.logger
+    @logger = StringIO.new
+    Lhm.logger = Logger.new(@logger)
+
+    @connection = DBConnectionHelper::new_mysql_connection(:proxysql, true, true)
+
+    @lhm_retry = Lhm::SqlRetry.new(@connection, retry_options: {},
+                                   reconnect_with_consistent_host: true)
+  end
+
+  after(:each) do
+    # Restore default logger
+    Lhm.logger = @old_logger
+  end
+
+  it "Will abort if service is down" do
+
+    e = assert_raises Lhm::Error do
+      #Service down
+      Toxiproxy[:mysql_proxysql].down do
+        @lhm_retry.with_retries do |retriable_connection|
+          retriable_connection.execute("INSERT INTO #{DBConnectionHelper.test_table_name} (id) VALUES (2000)")
+        end
+      end
+    end
+    assert_equal Lhm::Error, e.class
+    assert_match(/LHM tried the reconnection procedure but failed. Latest error:/, e.message)
+  end
+
+  it "Will retry until connection is achieved" do
+
+    #Creating a network blip
+    ToxiproxyHelper.with_kill_and_restart(:mysql_proxysql, 2.seconds) do
+      @lhm_retry.with_retries do |retriable_connection|
+        retriable_connection.execute("INSERT INTO #{DBConnectionHelper.test_table_name} (id) VALUES (2000)")
+      end
+    end
+
+    assert_equal @connection.execute("Select * from #{DBConnectionHelper.test_table_name} WHERE id=2000").to_a.first.first, 2000
+
+    logs = @logger.string.split("\n")
+
+    assert logs.first.include?("Lost connection to MySQL, will retry to connect to same host")
+    assert logs.last.include?("LHM successfully reconnected to initial host")
+  end
+
+  it "Will abort if new writer is not same host" do
+    # The hostname will be constant before the blip
+    Lhm::SqlRetry.any_instance.stubs(:hostname).returns("mysql-1").then.returns("mysql-2")
+
+    # Need new instance for stub to take into effect
+    lhm_retry = Lhm::SqlRetry.new(@connection, retry_options: {},
+                                  reconnect_with_consistent_host: true)
+
+    e = assert_raises Lhm::Error do
+      #Creating a network blip
+      ToxiproxyHelper.with_kill_and_restart(:mysql_proxysql, 2.seconds) do
+        lhm_retry.with_retries do |retriable_connection|
+          retriable_connection.execute("INSERT INTO #{DBConnectionHelper.test_table_name} (id) VALUES (2000)")
+        end
+      end
+    end
+
+    assert_equal e.class, Lhm::Error
+    assert_match(/LHM tried the reconnection procedure but failed. Latest error: Reconnected to wrong host/, e.message)
+
+    logs = @logger.string.split("\n")
+
+    assert logs.first.include?("Lost connection to MySQL, will retry to connect to same host")
+    assert logs.last.include?("Lost connection to MySQL server at 'reading initial communication packet")
+  end
+
+  it "Will abort if failover happens (mimicked with proxySQL)" do
+    e = assert_raises Lhm::Error do
+      #Creates a failover by switching the target hostgroup for the #hostname
+      ProxySQLHelper.with_lhm_hostgroup_flip do
+        #Creating a network blip
+        ToxiproxyHelper.with_kill_and_restart(:mysql_proxysql, 2.seconds) do
+          @lhm_retry.with_retries do |retriable_connection|
+            retriable_connection.execute("INSERT INTO #{DBConnectionHelper.test_table_name} (id) VALUES (2000)")
+          end
+        end
+      end
+    end
+
+    assert_equal e.class, Lhm::Error
+    assert_match(/LHM tried the reconnection procedure but failed. Latest error: Reconnected to wrong host/, e.message)
+
+    logs = @logger.string.split("\n")
+
+    assert logs.first.include?("Lost connection to MySQL, will retry to connect to same host")
+    assert logs.last.include?("Lost connection to MySQL server at 'reading initial communication packet")
+  end
+end

data/spec/integration/toxiproxy_helper.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+require 'toxiproxy'
+
+module ToxiproxyHelper
+  class << self
+
+    def included(base)
+      Toxiproxy.reset
+
+      # listen on localhost, but toxiproxy is in a container itself, thus the upstream uses the Docker-Compose DNS
+      Toxiproxy.populate(
+        [
+          {
+            name: 'mysql_master',
+            listen: '0.0.0.0:22220',
+            upstream: 'mysql-1:3306'
+          },
+          {
+            name: 'mysql_proxysql',
+            listen: '0.0.0.0:22222',
+            upstream: 'proxysql:3306'
+          }
+        ])
+    end
+
+    def with_kill_and_restart(target, restart_after)
+      thread = Thread.new do
+        sleep(restart_after) unless restart_after.nil?
+        Toxiproxy[target].enable
+      end
+
+      Toxiproxy[target].disable
+
+      yield
+
+    ensure
+      thread.join
+    end
+  end
+end

data/spec/test_helper.rb

@@ -10,6 +10,8 @@ require 'minitest/autorun'
 require 'minitest/spec'
 require 'minitest/mock'
 require 'mocha/minitest'
+require 'after_do'
+require 'byebug'
 require 'pathname'
 require 'lhm'
 
@@ -17,6 +19,8 @@ $project = Pathname.new(File.dirname(__FILE__) + '/..').cleanpath
 $spec = $project.join('spec')
 $fixtures = $spec.join('fixtures')
 
+$db_name = 'test'
+
 require 'active_record'
 require 'mysql2'
 
@@ -43,3 +47,20 @@ end
 def throttler
   Lhm::Throttler::Time.new(:stride => 100)
 end
+
+def init_test_db
+  db_config = YAML.load_file(File.expand_path(File.dirname(__FILE__)) + '/integration/database.yml')
+  conn = Mysql2::Client.new(
+    :host => '127.0.0.1',
+    :username => db_config['master']['user'],
+    :password => db_config['master']['password'],
+    :port => db_config['master']['port']
+  )
+
+  conn.query("DROP DATABASE IF EXISTS #{$db_name}")
+  conn.query("CREATE DATABASE #{$db_name}")
+end
+
+init_test_db
+
+

data/spec/unit/chunk_insert_spec.rb

@@ -4,17 +4,22 @@
 require File.expand_path(File.dirname(__FILE__)) + '/unit_helper'
 
 require 'lhm/chunk_insert'
+require 'lhm/connection'
 
 describe Lhm::ChunkInsert do
   before(:each) do
-    @connection = stub(:connection)
+    ar_connection = mock()
+    ar_connection.stubs(:execute).returns([["dummy"]])
+    @connection = Lhm::Connection.new(connection: ar_connection, options: {reconnect_with_consistent_host: false})
     @origin = Lhm::Table.new('foo')
     @destination = Lhm::Table.new('bar')
   end
 
   describe "#sql" do
     describe "when migration has no conditions" do
-      before { @migration = Lhm::Migration.new(@origin, @destination) }
+      before do
+        @migration = Lhm::Migration.new(@origin, @destination)
+      end
 
       it "uses a simple where clause" do
         assert_equal(

data/spec/unit/chunker_spec.rb

@@ -7,15 +7,19 @@ require 'lhm/table'
 require 'lhm/migration'
 require 'lhm/chunker'
 require 'lhm/throttler'
+require 'lhm/connection'
 
 describe Lhm::Chunker do
   include UnitHelper
 
+  EXPECTED_RETRY_FLAGS = {:should_retry => true, :retry_options => {}}
+
   before(:each) do
     @origin = Lhm::Table.new('foo')
     @destination = Lhm::Table.new('bar')
     @migration = Lhm::Migration.new(@origin, @destination)
     @connection = mock()
+    @connection.stubs(:execute).returns([["dummy"]])
     # This is a poor man's stub
     @throttler = Object.new
     def @throttler.run
@@ -37,11 +41,11 @@ describe Lhm::Chunker do
         5
       end
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 4/)).returns(7)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 8 order by id limit 1 offset 4/)).returns(21)
-      @connection.expects(:update).with(regexp_matches(/between 1 and 7/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 8 and 10/)).returns(2)
-      @connection.expects(:query).twice.with(regexp_matches(/show warnings/)).returns([])
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 4/),EXPECTED_RETRY_FLAGS).returns(7)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 8 order by id limit 1 offset 4/),EXPECTED_RETRY_FLAGS).returns(21)
+      @connection.expects(:update).with(regexp_matches(/between 1 and 7/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 8 and 10/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:execute).twice.with(regexp_matches(/show warnings/),EXPECTED_RETRY_FLAGS).returns([])
 
       @chunker.run
     end
@@ -52,17 +56,17 @@ describe Lhm::Chunker do
         2
       end
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/)).returns(2)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 3 order by id limit 1 offset 1/)).returns(4)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 5 order by id limit 1 offset 1/)).returns(6)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 7 order by id limit 1 offset 1/)).returns(8)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 9 order by id limit 1 offset 1/)).returns(10)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 3 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(4)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 5 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(6)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 7 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(8)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 9 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(10)
 
-      @connection.expects(:update).with(regexp_matches(/between 1 and 2/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 3 and 4/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 5 and 6/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 7 and 8/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 9 and 10/)).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 1 and 2/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 3 and 4/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 5 and 6/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 7 and 8/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 9 and 10/),EXPECTED_RETRY_FLAGS).returns(2)
 
       @chunker.run
     end
@@ -79,17 +83,17 @@ describe Lhm::Chunker do
         end
       end
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/)).returns(2)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 3 order by id limit 1 offset 2/)).returns(5)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 6 order by id limit 1 offset 2/)).returns(8)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 9 order by id limit 1 offset 2/)).returns(nil)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 3 order by id limit 1 offset 2/),EXPECTED_RETRY_FLAGS).returns(5)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 6 order by id limit 1 offset 2/),EXPECTED_RETRY_FLAGS).returns(8)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 9 order by id limit 1 offset 2/),EXPECTED_RETRY_FLAGS).returns(nil)
 
-      @connection.expects(:update).with(regexp_matches(/between 1 and 2/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 3 and 5/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 6 and 8/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 9 and 10/)).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 1 and 2/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 3 and 5/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 6 and 8/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 9 and 10/),EXPECTED_RETRY_FLAGS).returns(2)
 
-      @connection.expects(:query).twice.with(regexp_matches(/show warnings/)).returns([])
+      @connection.expects(:execute).twice.with(regexp_matches(/show warnings/),EXPECTED_RETRY_FLAGS).returns([])
 
       @chunker.run
     end
@@ -99,8 +103,8 @@ describe Lhm::Chunker do
                                                            :start     => 1,
                                                            :limit     => 1)
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 0/)).returns(nil)
-      @connection.expects(:update).with(regexp_matches(/between 1 and 1/)).returns(1)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 0/),EXPECTED_RETRY_FLAGS).returns(nil)
+      @connection.expects(:update).with(regexp_matches(/between 1 and 1/),EXPECTED_RETRY_FLAGS).returns(1)
 
       @chunker.run
     end
@@ -113,17 +117,17 @@ describe Lhm::Chunker do
         2
       end
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 2 order by id limit 1 offset 1/)).returns(3)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 4 order by id limit 1 offset 1/)).returns(5)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 6 order by id limit 1 offset 1/)).returns(7)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 8 order by id limit 1 offset 1/)).returns(9)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 10 order by id limit 1 offset 1/)).returns(nil)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 2 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(3)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 4 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(5)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 6 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(7)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 8 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(9)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 10 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(nil)
 
-      @connection.expects(:update).with(regexp_matches(/between 2 and 3/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 4 and 5/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 6 and 7/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 8 and 9/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 10 and 10/)).returns(1)
+      @connection.expects(:update).with(regexp_matches(/between 2 and 3/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 4 and 5/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 6 and 7/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 8 and 9/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 10 and 10/),EXPECTED_RETRY_FLAGS).returns(1)
 
       @chunker.run
     end
@@ -137,9 +141,9 @@ describe Lhm::Chunker do
         2
       end
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/where \(foo.created_at > '2013-07-10' or foo.baz = 'quux'\) and `foo`/)).returns(1)
-      @connection.expects(:query).with(regexp_matches(/show warnings/)).returns([])
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/where \(foo.created_at > '2013-07-10' or foo.baz = 'quux'\) and `foo`/),EXPECTED_RETRY_FLAGS).returns(1)
+      @connection.expects(:execute).with(regexp_matches(/show warnings/),EXPECTED_RETRY_FLAGS).returns([])
 
       def @migration.conditions
         "where foo.created_at > '2013-07-10' or foo.baz = 'quux'"
@@ -157,9 +161,9 @@ describe Lhm::Chunker do
         2
       end
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/)).returns(2)
-      @connection.expects(:update).with(regexp_matches(/inner join bar on foo.id = bar.foo_id and/)).returns(1)
-      @connection.expects(:query).with(regexp_matches(/show warnings/)).returns([])
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS).returns(2)
+      @connection.expects(:update).with(regexp_matches(/inner join bar on foo.id = bar.foo_id and/),EXPECTED_RETRY_FLAGS).returns(1)
+      @connection.expects(:execute).with(regexp_matches(/show warnings/),EXPECTED_RETRY_FLAGS).returns([])
 
       def @migration.conditions
         'inner join bar on foo.id = bar.foo_id'

data/spec/unit/connection_spec.rb

@@ -0,0 +1,86 @@
+require 'lhm/connection'
+require 'lhm/proxysql_helper'
+
+describe Lhm::Connection do
+
+  LOCK_WAIT = ActiveRecord::StatementInvalid.new('Lock wait timeout exceeded; try restarting transaction.')
+
+  before(:each) do
+    @logs = StringIO.new
+    Lhm.logger = Logger.new(@logs)
+  end
+
+  it "Should find use calling file as prefix" do
+    ar_connection = mock()
+    ar_connection.stubs(:execute).raises(LOCK_WAIT).then.returns(true)
+    ar_connection.stubs(:active?).returns(true)
+
+    connection = Lhm::Connection.new(connection: ar_connection)
+
+    connection.execute("SHOW TABLES", should_retry: true, retry_options: { base_interval: 0 })
+
+    log_messages = @logs.string.split("\n")
+    assert_equal(1, log_messages.length)
+    assert log_messages.first.include?("[ConnectionSpec]")
+  end
+
+  it "#execute should be retried" do
+    ar_connection = mock()
+    ar_connection.stubs(:execute).raises(LOCK_WAIT)
+                 .then.raises(LOCK_WAIT)
+                 .then.returns(true)
+    ar_connection.stubs(:active?).returns(true)
+
+    connection = Lhm::Connection.new(connection: ar_connection)
+
+    connection.execute("SHOW TABLES", should_retry: true, retry_options: { base_interval: 0, tries: 3 })
+
+    log_messages = @logs.string.split("\n")
+    assert_equal(2, log_messages.length)
+  end
+
+  it "#update should be retried" do
+    ar_connection = mock()
+    ar_connection.stubs(:update).raises(LOCK_WAIT)
+                 .then.raises(LOCK_WAIT)
+                 .then.returns(1)
+    ar_connection.stubs(:active?).returns(true)
+
+    connection = Lhm::Connection.new(connection: ar_connection)
+
+    val = connection.update("SHOW TABLES", should_retry: true, retry_options:{ base_interval: 0, tries: 3 })
+
+    log_messages = @logs.string.split("\n")
+    assert_equal val, 1
+    assert_equal(2, log_messages.length)
+  end
+
+  it "#select_value should be retried" do
+    ar_connection = mock()
+    ar_connection.stubs(:select_value).raises(LOCK_WAIT)
+                 .then.raises(LOCK_WAIT)
+                 .then.returns("dummy")
+    ar_connection.stubs(:active?).returns(true)
+
+    connection = Lhm::Connection.new(connection: ar_connection)
+
+    val = connection.select_value("SHOW TABLES", should_retry: true, retry_options: { base_interval: 0, tries: 3 })
+
+    log_messages = @logs.string.split("\n")
+    assert_equal val, "dummy"
+    assert_equal(2, log_messages.length)
+  end
+
+  it "Queries should be tagged with ProxySQL tag if requested" do
+    ar_connection = mock()
+    ar_connection.expects(:public_send).with(:select_value, "SHOW TABLES #{Lhm::ProxySQLHelper::ANNOTATION}").returns("dummy")
+    ar_connection.stubs(:execute).times(4).returns([["dummy"]])
+    ar_connection.stubs(:active?).returns(true)
+
+    connection = Lhm::Connection.new(connection: ar_connection, options: { reconnect_with_consistent_host: true })
+
+    val = connection.select_value("SHOW TABLES", should_retry: true, retry_options: { base_interval: 0, tries: 3 })
+
+    assert_equal val, "dummy"
+  end
+end

data/spec/unit/entangler_spec.rb

@@ -6,6 +6,7 @@ require File.expand_path(File.dirname(__FILE__)) + '/unit_helper'
 require 'lhm/table'
 require 'lhm/migration'
 require 'lhm/entangler'
+require 'lhm/connection'
 
 describe Lhm::Entangler do
   include UnitHelper
@@ -60,33 +61,73 @@ describe Lhm::Entangler do
     end
 
     it 'should retry trigger creation when it hits a lock wait timeout' do
-      connection = mock()
       tries = 1
+      ar_connection = mock()
+      ar_connection.stubs(:execute)
+                       .returns([["dummy"]], [["dummy"]], [["dummy"]])
+                       .then
+                       .raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
+      ar_connection.stubs(:active?).returns(true)
+
+      connection = Lhm::Connection.new(connection: ar_connection, options: {reconnect_with_consistent_host: true})
+
       @entangler = Lhm::Entangler.new(@migration, connection, retriable: {base_interval: 0, tries: tries})
-      connection.expects(:execute).times(tries).raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
 
       assert_raises(Mysql2::Error) { @entangler.before }
     end
 
     it 'should not retry trigger creation with other mysql errors' do
-      connection = mock()
-      connection.expects(:execute).once.raises(Mysql2::Error, 'The MySQL server is running with the --read-only option so it cannot execute this statement.')
-
-      @entangler = Lhm::Entangler.new(@migration, connection, retriable: {base_interval: 0})
+      ar_connection = mock()
+      ar_connection.stubs(:execute)
+                   .returns([["dummy"]], [["dummy"]], [["dummy"]])
+                   .then
+                   .raises(Mysql2::Error, 'The MySQL server is running with the --read-only option so it cannot execute this statement.')
+      ar_connection.stubs(:active?).returns(true)
+      connection = Lhm::Connection.new(connection: ar_connection, options: {reconnect_with_consistent_host: true})
+
+      @entangler = Lhm::Entangler.new(@migration, connection, retriable: { base_interval: 0 })
       assert_raises(Mysql2::Error) { @entangler.before }
     end
 
     it 'should succesfully finish after retrying' do
-      connection = mock()
-      connection.stubs(:execute).raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction').then.returns(true)
+      ar_connection = mock()
+      ar_connection.stubs(:execute)
+                   .returns([["dummy"]], [["dummy"]], [["dummy"]])
+                   .then
+                   .raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
+                   .then
+                   .returns([["dummy"]])
+      ar_connection.stubs(:active?).returns(true)
+
+      connection = Lhm::Connection.new(connection: ar_connection, options: {reconnect_with_consistent_host: true})
+
       @entangler = Lhm::Entangler.new(@migration, connection, retriable: {base_interval: 0})
 
       assert @entangler.before
     end
 
     it 'should retry as many times as specified by configuration' do
-      connection = mock()
-      connection.expects(:execute).times(5).raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
+      ar_connection = mock()
+      ar_connection.stubs(:execute)
+                   .returns([["dummy"]], [["dummy"]], [["dummy"]]) # initial
+                   .then
+                   .raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
+                   .then
+                   .returns([["dummy"]]) # reconnect 1
+                   .then
+                   .raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
+                   .then
+                   .returns([["dummy"]])  # reconnect 2
+                   .then
+                   .raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
+                   .then
+                   .returns([["dummy"]])  # reconnect 3
+                   .then
+                   .raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')  # final error
+      ar_connection.stubs(:active?).returns(true)
+
+      connection = Lhm::Connection.new(connection: ar_connection, options: {reconnect_with_consistent_host: true})
+
       @entangler = Lhm::Entangler.new(@migration, connection, retriable: {tries: 5, base_interval: 0})
 
       assert_raises(Mysql2::Error) { @entangler.before }

data/spec/unit/lhm_spec.rb

@@ -26,4 +26,21 @@ describe Lhm do
       value(Lhm.logger.instance_eval { @logdev }.dev.path).must_equal 'omg.ponies'
     end
   end
+
+  describe 'api' do
+
+    before(:each) do
+      @connection = mock()
+    end
+
+    it 'should create a new connection when calling setup' do
+      Lhm.setup(@connection)
+      value(Lhm.connection).must_be_kind_of(Lhm::Connection)
+    end
+
+    it 'should create a new connection when none is created' do
+      ActiveRecord::Base.stubs(:connection).returns(@connection)
+      value(Lhm.connection).must_be_kind_of(Lhm::Connection)
+    end
+  end
 end

data/spec/unit/throttler/slave_lag_spec.rb

@@ -39,7 +39,7 @@ describe Lhm::Throttler::Slave do
     @logs = StringIO.new
     Lhm.logger = Logger.new(@logs)
 
-    @dummy_mysql_client_config = lambda { {'username' => 'user', 'password' => 'pw', 'database' => 'db'} }
+    @dummy_mysql_client_config = lambda { { 'username' => 'user', 'password' => 'pw', 'database' => 'db' } }
   end
 
   describe "#client" do
@@ -64,7 +64,7 @@ describe Lhm::Throttler::Slave do
 
     describe 'with proper config' do
       it "creates a new Mysql2::Client" do
-        expected_config = {username: 'user', password: 'pw', database: 'db', host: 'slave'}
+        expected_config = { username: 'user', password: 'pw', database: 'db', host: 'slave' }
         Mysql2::Client.stubs(:new).with(expected_config).returns(mock())
 
         assert Lhm::Throttler::Slave.new('slave', @dummy_mysql_client_config).connection
@@ -73,8 +73,12 @@ describe Lhm::Throttler::Slave do
 
     describe 'with active record config' do
       it 'logs and creates client' do
-        active_record_config = {username: 'user', password: 'pw', database: 'db'}
-        ActiveRecord::Base.stubs(:connection_pool).returns(stub(spec: stub(config: active_record_config)))
+        active_record_config = { username: 'user', password: 'pw', database: 'db' }
+        if ActiveRecord::VERSION::MAJOR > 6 || ActiveRecord::VERSION::MAJOR == 6 && ActiveRecord::VERSION::MINOR >= 1
+          ActiveRecord::Base.stubs(:connection_pool).returns(stub(db_config: stub(configuration_hash: active_record_config)))
+        else
+          ActiveRecord::Base.stubs(:connection_pool).returns(stub(spec: stub(config: active_record_config)))
+        end
 
         Mysql2::Client.stubs(:new).returns(mock())
 
@@ -92,9 +96,9 @@ describe Lhm::Throttler::Slave do
       class Connection
         def self.query(query)
           if query == Lhm::Throttler::Slave::SQL_SELECT_MAX_SLAVE_LAG
-            [{'Seconds_Behind_Master' => 20}]
+            [{ 'Seconds_Behind_Master' => 20 }]
           elsif query == Lhm::Throttler::Slave::SQL_SELECT_SLAVE_HOSTS
-            [{'host' => '1.1.1.1:80'}]
+            [{ 'host' => '1.1.1.1:80' }]
           end
         end
       end
@@ -104,7 +108,7 @@ describe Lhm::Throttler::Slave do
 
       class StoppedConnection
         def self.query(query)
-          [{'Seconds_Behind_Master' => nil}]
+          [{ 'Seconds_Behind_Master' => nil }]
         end
       end
 
@@ -286,7 +290,7 @@ describe Lhm::Throttler::SlaveLag do
       describe 'with the :check_only option' do
         describe 'with a callable argument' do
           before do
-            check_only = lambda {{'host' => '1.1.1.3'}}
+            check_only = lambda { { 'host' => '1.1.1.3' } }
             @throttler = Lhm::Throttler::SlaveLag.new :check_only => check_only
           end
 
@@ -300,6 +304,7 @@ describe Lhm::Throttler::SlaveLag do
         describe 'with a non-callable argument' do
           before do
             @throttler = Lhm::Throttler::SlaveLag.new :check_only => 'I cannot be called'
+
             def @throttler.master_slave_hosts
               ['1.1.1.1', '1.1.1.4']
             end