Fixing rollback issues in Rails 4

Gemfile.lock

@@ -1,34 +1,40 @@
 PATH
   remote: .
   specs:
-    mysql_online_migrations (0.1.4)
-      activerecord (~> 3.2.15)
-      activesupport (~> 3.2.15)
+    mysql_online_migrations (1.0.1)
+      activerecord (>= 3.2.15)
+      activesupport (>= 3.2.15)
       mysql2
 
 GEM
   remote: http://rubygems.org/
   specs:
-    activemodel (3.2.15)
-      activesupport (= 3.2.15)
-      builder (~> 3.0.0)
-    activerecord (3.2.15)
-      activemodel (= 3.2.15)
-      activesupport (= 3.2.15)
-      arel (~> 3.0.2)
-      tzinfo (~> 0.3.29)
-    activesupport (3.2.15)
+    activemodel (4.0.2)
+      activesupport (= 4.0.2)
+      builder (~> 3.1.0)
+    activerecord (4.0.2)
+      activemodel (= 4.0.2)
+      activerecord-deprecated_finders (~> 1.0.2)
+      activesupport (= 4.0.2)
+      arel (~> 4.0.0)
+    activerecord-deprecated_finders (1.0.3)
+    activesupport (4.0.2)
       i18n (~> 0.6, >= 0.6.4)
-      multi_json (~> 1.0)
-    arel (3.0.3)
-    builder (3.0.4)
+      minitest (~> 4.2)
+      multi_json (~> 1.3)
+      thread_safe (~> 0.1)
+      tzinfo (~> 0.3.37)
+    arel (4.0.2)
+    atomic (1.1.14)
+    builder (3.1.4)
     coderay (1.0.9)
     diff-lcs (1.2.5)
     i18n (0.6.9)
     logger (1.2.8)
     method_source (0.8.2)
-    multi_json (1.8.2)
-    mysql2 (0.3.14)
+    minitest (4.7.5)
+    multi_json (1.8.4)
+    mysql2 (0.3.15)
     pry (0.9.12.2)
       coderay (~> 1.0.5)
       method_source (~> 0.8)
@@ -42,6 +48,8 @@ GEM
       diff-lcs (>= 1.1.3, < 2.0)
     rspec-mocks (2.14.4)
     slop (3.4.6)
+    thread_safe (0.1.3)
+      atomic
     tzinfo (0.3.38)
 
 PLATFORMS

lib/mysql_online_migrations.rb

@@ -13,7 +13,23 @@ def self.prepended(base)
   end
 
   def connection
-    @no_lock_adapter ||= ActiveRecord::ConnectionAdapters::Mysql2AdapterWithoutLock.new(super)
+    original_connection = super
+    adapter_mode = original_connection.class.name == "ActiveRecord::ConnectionAdapters::Mysql2Adapter"
+
+    @original_adapter ||= if adapter_mode
+      original_connection
+    else
+      original_connection.instance_variable_get(:@delegate)
+    end
+
+    @no_lock_adapter ||= ActiveRecord::ConnectionAdapters::Mysql2AdapterWithoutLock.new(@original_adapter)
+
+    if adapter_mode
+      @no_lock_adapter
+    else
+      original_connection.instance_variable_set(:@delegate, @no_lock_adapter)
+      original_connection
+    end
   end
 
   def with_lock

mysql_online_migrations.gemspec

@@ -1,6 +1,6 @@
 Gem::Specification.new do |s|
   s.name        = 'mysql_online_migrations'
-  s.version     = '1.0.1'
+  s.version     = '1.0.2'
   s.summary     = "Use MySQL 5.6+ capacities to enforce online migrations"
   s.description = "MySQL 5.6 adds a `LOCK=NONE` option to make sure migrations are done with no locking. Let's use it."
   s.authors     = ["Anthony Alberto"]

spec/lib/mysql_online_migrations_spec.rb

@@ -10,11 +10,34 @@
   end
 
   context "#connection" do
-    it "memoizes an instance of Mysql2AdapterWithoutLock" do
-      ActiveRecord::ConnectionAdapters::Mysql2AdapterWithoutLock.should_receive(:new)
-        .with(an_instance_of(ActiveRecord::ConnectionAdapters::Mysql2Adapter)).once.and_call_original
-      3.times { @connection = migration.connection }
-      @connection.should be_an_instance_of(ActiveRecord::ConnectionAdapters::Mysql2AdapterWithoutLock)
+    shared_examples_for "Mysql2AdapterWithoutLock created" do
+      it "memoizes an instance of Mysql2AdapterWithoutLock" do
+        ActiveRecord::ConnectionAdapters::Mysql2AdapterWithoutLock.should_receive(:new)
+          .with(an_instance_of(ActiveRecord::ConnectionAdapters::Mysql2Adapter)).once.and_call_original
+        3.times { migration.connection }
+      end
+    end
+
+    context 'when migrating' do
+      it "returns an instance of Mysql2AdapterWithoutLock" do
+        migration.connection.should be_an_instance_of(ActiveRecord::ConnectionAdapters::Mysql2AdapterWithoutLock)
+      end
+
+      it_behaves_like "Mysql2AdapterWithoutLock created"
+    end
+
+    context 'when rolling back' do
+      before do
+        migration.instance_variable_set(:@connection, ActiveRecord::Migration::CommandRecorder.new(ActiveRecord::Base.connection))
+      end
+
+      it "returns an instance of ActiveRecord::Migration::CommandRecorder" do
+        recorder_connection = migration.connection
+        recorder_connection.should be_an_instance_of(ActiveRecord::Migration::CommandRecorder)
+        recorder_connection.instance_variable_get(:@delegate).should be_an_instance_of(ActiveRecord::ConnectionAdapters::Mysql2AdapterWithoutLock)
+      end
+
+      it_behaves_like "Mysql2AdapterWithoutLock created"
     end
   end