lhm-shopify 3.4.0 → 3.5.5

data/spec/integration/sql_retry/retry_with_proxysql_spec.rb

@@ -0,0 +1,109 @@
+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. Aborting/, 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")
+    Lhm::SqlRetry.any_instance.stubs(:server_id).returns(1).then.returns(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. Aborting/, 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?("Reconnected to wrong host. Started migration on: mysql-1 (server_id: 1), but reconnected to: mysql-2 (server_id: 2).")
+  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. Aborting/, 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?("Reconnected to wrong host. Started migration on: mysql-1 (server_id: 1), but reconnected to: mysql-2 (server_id: 2).")
+  end
+end

data/spec/integration/table_spec.rb

@@ -15,28 +15,24 @@ describe Lhm::Table do
       end
 
       it 'should parse primary key' do
-        @table.pk.must_equal('pk')
+        value(@table.pk).must_equal('pk')
       end
 
       it 'should parse indices' do
-        @table.
-          indices['index_custom_primary_key_on_id'].
-          must_equal(['id'])
+        value(@table.indices['index_custom_primary_key_on_id']).must_equal(['id'])
       end
 
       it 'should parse columns' do
-        @table.
-          columns['id'][:type].
-          must_match(/(bigint|int)\(\d+\)/)
+        value(@table.columns['id'][:type]).must_match(/(bigint|int)\(\d+\)/)
       end
 
       it 'should return true for method that should be renamed' do
-        @table.satisfies_id_column_requirement?.must_equal true
+        value(@table.satisfies_id_column_requirement?).must_equal true
       end
 
       it 'should support bigint tables' do
         @table = table_create(:bigint_table)
-        @table.satisfies_id_column_requirement?.must_equal true
+        value(@table.satisfies_id_column_requirement?).must_equal true
       end
     end
 
@@ -47,7 +43,7 @@ describe Lhm::Table do
 
       it 'should return false for a non-int id column' do
         @table = table_create(:wo_id_int_column)
-        @table.satisfies_id_column_requirement?.must_equal false
+        value(@table.satisfies_id_column_requirement?).must_equal false
       end
     end
   end
@@ -60,15 +56,15 @@ describe Lhm::Table do
       end
 
       it 'should parse table name in show create table' do
-        @table.name.must_equal('users')
+        value(@table.name).must_equal('users')
       end
 
       it 'should parse primary key' do
-        @table.pk.must_equal('id')
+        value(@table.pk).must_equal('id')
       end
 
       it 'should parse column type in show create table' do
-        @table.columns['username'][:type].must_equal('varchar(255)')
+        value(@table.columns['username'][:type]).must_equal('varchar(255)')
       end
 
       it 'should parse column metadata' do
@@ -76,15 +72,11 @@ describe Lhm::Table do
       end
 
       it 'should parse indices' do
-        @table.
-          indices['index_users_on_username_and_created_at'].
-          must_equal(['username', 'created_at'])
+        value(@table.indices['index_users_on_username_and_created_at']).must_equal(['username', 'created_at'])
       end
 
       it 'should parse index' do
-        @table.
-          indices['index_users_on_reference'].
-          must_equal(['reference'])
+        value(@table.indices['index_users_on_reference']).must_equal(['reference'])
       end
     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'
 
@@ -24,6 +28,9 @@ logger = Logger.new STDOUT
 logger.level = Logger::WARN
 Lhm.logger = logger
 
+# Want test to be efficient without having to wait the normal value of 120s
+Lhm::SqlRetry::RECONNECT_RETRY_MAX_ITERATION = 4
+
 def without_verbose(&block)
   old_verbose, $VERBOSE = $VERBOSE, nil
   yield
@@ -43,3 +50,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/atomic_switcher_spec.rb

@@ -20,12 +20,10 @@ describe Lhm::AtomicSwitcher do
 
   describe 'atomic switch' do
     it 'should perform a single atomic rename' do
-      @switcher.
-        atomic_switch.
-        must_equal(
-          "rename table `origin` to `#{ @migration.archive_name }`, " \
-          '`destination` to `origin`'
-        )
+      value(@switcher.atomic_switch).must_equal(
+        "rename table `origin` to `#{ @migration.archive_name }`, " \
+        '`destination` to `origin`'
+      )
     end
   end
 end

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,20 @@ require 'lhm/table'
 require 'lhm/migration'
 require 'lhm/chunker'
 require 'lhm/throttler'
+require 'lhm/connection'
 
 describe Lhm::Chunker do
   include UnitHelper
 
+  EXPECTED_RETRY_FLAGS_CHUNKER = {:should_retry => true, :log_prefix => "Chunker"}
+  EXPECTED_RETRY_FLAGS_CHUNK_INSERT = {:should_retry => true, :log_prefix => "ChunkInsert"}
+
   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 +42,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_CHUNKER).returns(7)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 8 order by id limit 1 offset 4/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(21)
+      @connection.expects(:update).with(regexp_matches(/between 1 and 7/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 8 and 10/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:execute).twice.with(regexp_matches(/show warnings/),EXPECTED_RETRY_FLAGS_CHUNKER).returns([])
 
       @chunker.run
     end
@@ -52,17 +57,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_CHUNKER).returns(2)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 3 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(4)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 5 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(6)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 7 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(8)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 9 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).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_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 3 and 4/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 5 and 6/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 7 and 8/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 9 and 10/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
 
       @chunker.run
     end
@@ -79,17 +84,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_CHUNKER).returns(2)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 3 order by id limit 1 offset 2/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(5)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 6 order by id limit 1 offset 2/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(8)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 9 order by id limit 1 offset 2/),EXPECTED_RETRY_FLAGS_CHUNKER).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_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 3 and 5/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 6 and 8/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 9 and 10/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
 
-      @connection.expects(:query).twice.with(regexp_matches(/show warnings/)).returns([])
+      @connection.expects(:execute).twice.with(regexp_matches(/show warnings/),EXPECTED_RETRY_FLAGS_CHUNKER).returns([])
 
       @chunker.run
     end
@@ -99,8 +104,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_CHUNKER).returns(nil)
+      @connection.expects(:update).with(regexp_matches(/between 1 and 1/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(1)
 
       @chunker.run
     end
@@ -113,17 +118,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_CHUNKER).returns(3)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 4 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(5)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 6 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(7)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 8 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(9)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 10 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).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_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 4 and 5/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 6 and 7/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 8 and 9/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 10 and 10/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(1)
 
       @chunker.run
     end
@@ -137,9 +142,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_CHUNKER).returns(2)
+      @connection.expects(:update).with(regexp_matches(/where \(foo.created_at > '2013-07-10' or foo.baz = 'quux'\) and `foo`/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(1)
+      @connection.expects(:execute).with(regexp_matches(/show warnings/),EXPECTED_RETRY_FLAGS_CHUNKER).returns([])
 
       def @migration.conditions
         "where foo.created_at > '2013-07-10' or foo.baz = 'quux'"
@@ -157,9 +162,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_CHUNKER).returns(2)
+      @connection.expects(:update).with(regexp_matches(/inner join bar on foo.id = bar.foo_id and/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(1)
+      @connection.expects(:execute).with(regexp_matches(/show warnings/),EXPECTED_RETRY_FLAGS_CHUNKER).returns([])
 
       def @migration.conditions
         'inner join bar on foo.id = bar.foo_id'

data/spec/unit/connection_spec.rb

@@ -0,0 +1,111 @@
+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, options: {
+      retriable: {
+        base_interval: 0
+      }
+    })
+
+    connection.execute("SHOW TABLES", should_retry: true)
+
+    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, options: {
+      retriable: {
+        base_interval: 0,
+        tries: 3
+      }
+    })
+
+    connection.execute("SHOW TABLES", should_retry: true)
+
+    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, options: {
+      retriable: {
+        base_interval: 0,
+        tries: 3
+      }
+    })
+
+    val = connection.update("SHOW TABLES", should_retry: true)
+
+    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, options: {
+      retriable: {
+        base_interval: 0,
+        tries: 3
+      }
+    })
+
+    val = connection.select_value("SHOW TABLES", should_retry: true)
+
+    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 reconnect_with_consistent_host is enabled" 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,
+      retriable: {
+        base_interval: 0,
+        tries: 3
+      }
+    })
+
+    val = connection.select_value("SHOW TABLES", should_retry: true)
+
+    assert_equal val, "dummy"
+  end
+end