Merge pull request #51 from gussan/feature/rails51

Support rails 5.1

.travis.yml

@@ -15,6 +15,8 @@ gemfile:
   - gemfiles/rails5_0_1.gemfile
   - gemfiles/rails5_0_2.gemfile
   - gemfiles/rails5_0_3.gemfile
+  - gemfiles/rails5_1_0.gemfile
+  - gemfiles/rails5_1_1.gemfile
   - gemfiles/rails_edge.gemfile
 
 env:
@@ -40,7 +42,7 @@ matrix:
     - rvm: 2.4.1
       gemfile: gemfiles/rails_edge.gemfile
     - rvm: 2.3.4
-      gemfile: gemfiles/rails5_0_3.gemfile
+      gemfile: gemfiles/rails5_1_1.gemfile
   allow_failures:
     - rvm: ruby-head
     - gemfile: gemfiles/rails_edge.gemfile

Gemfile

@@ -2,5 +2,5 @@ source "https://rubygems.org"
 
 gemspec
 
-gem "activerecord", "~> 5.0.1"
-gem "activesupport", "~> 5.0.1"
+gem "activerecord", "~> 5.1.1"
+gem "activesupport", "~> 5.1.1"

Rakefile

@@ -129,7 +129,8 @@ namespace :turntable do
   namespace :activerecord do
     task(:env) do
       ENV["ARCONFIG"] ||= File.expand_path("spec/config/activerecord_config.yml", __dir__)
-      ENV["ARVERSION"] ||= if ActiveRecord.gem_version.prerelease?
+      ENV["ARVERSION"] ||= if ActiveRecord.gem_version.prerelease? &&
+                              !ActiveRecord.gem_version.segments.include?("rc")
                              "origin/master"
                            else
                              "v#{ActiveRecord.gem_version}"
@@ -143,6 +144,7 @@ namespace :turntable do
         Dir.chdir("tmp/rails") do
           system(*%W|git checkout #{ENV['ARVERSION']}|)
         end
+        FileUtils.rm_r("test") if File.directory?("test")
         FileUtils.cp_r("tmp/rails/activerecord/test", ".")
         FileUtils.cp_r("tmp/rails/activerecord/Rakefile", "activerecord.rake")
         File.open("test/cases/helper.rb", "a") do |f|

gemfiles/rails5_1_0.gemfile

@@ -0,0 +1,7 @@
+source "https://rubygems.org"
+
+gem "activerecord", "5.1.0"
+gem "activesupport", "5.1.0"
+gem "actionview", "5.1.0"
+
+gemspec :path => '../'

gemfiles/rails5_1_1.gemfile

@@ -0,0 +1,7 @@
+source "https://rubygems.org"
+
+gem "activerecord", "5.1.1"
+gem "activesupport", "5.1.1"
+gem "actionview", "5.1.1"
+
+gemspec :path => '../'

lib/active_record/turntable.rb

@@ -24,6 +24,7 @@ module ActiveRecord::Turntable
     autoload :ClusterHelperMethods
     autoload :Config
     autoload :ConnectionProxy
+    autoload :Compatibility
     autoload :MasterShard
     autoload :Migration
     autoload :Mixer

lib/active_record/turntable/active_record_ext/abstract_adapter.rb

@@ -1,6 +1,13 @@
 module ActiveRecord::Turntable
   module ActiveRecordExt
     module AbstractAdapter
+      extend Compatibility
+
+      def self.prepended(klass)
+        klass.prepend(self.compatible_module)
+        klass.class_eval { protected :log }
+      end
+
       def translate_exception_class(e, sql)
         begin
           message = "#{e.class.name}: #{e.message}: #{sql} : #{turntable_shard_name}"
@@ -12,11 +19,32 @@ def translate_exception_class(e, sql)
         exception.set_backtrace e.backtrace
         exception
       end
+      protected :translate_exception_class
 
       # @note override for append current shard name
       # rubocop:disable Style/HashSyntax, Style/MultilineMethodCallBraceLayout
-      if ActiveRecord::Turntable::Util.ar_version_equals_or_later?("5.0.3")
+      module V5_1
         def log(sql, name = "SQL", binds = [], type_casted_binds = [], statement_name = nil)
+          @instrumenter.instrument(
+            "sql.active_record",
+            sql:                  sql,
+            name:                 name,
+            binds:                binds,
+            type_casted_binds:    type_casted_binds,
+            statement_name:       statement_name,
+            connection_id:        object_id,
+            turntable_shard_name: turntable_shard_name) do
+              @lock.synchronize do
+                yield
+              end
+            end
+        rescue => e
+          raise translate_exception_class(e, sql)
+        end
+      end
+
+      module V5_0_3
+        def log(sql, name = "SQL", binds = [], type_casted_binds = [], statement_name = nil) # :doc:
           @instrumenter.instrument(
             "sql.active_record",
             sql:                  sql,
@@ -29,7 +57,9 @@ def log(sql, name = "SQL", binds = [], type_casted_binds = [], statement_name =
         rescue => e
           raise translate_exception_class(e, sql)
         end
-      else
+      end
+
+      module V5_0
         def log(sql, name = "SQL", binds = [], statement_name = nil)
           @instrumenter.instrument(
             "sql.active_record",
@@ -45,7 +75,6 @@ def log(sql, name = "SQL", binds = [], statement_name = nil)
       end
       # rubocop:enable Style/HashSyntax, Style/MultilineMethodCallBraceLayout
 
-      protected :translate_exception_class, :log
 
       def turntable_shard_name=(name)
         @turntable_shard_name = name.to_s

lib/active_record/turntable/active_record_ext/connection_handler_extension.rb

@@ -11,7 +11,9 @@ def retrieve_connection_pool(spec_name)
               # A connection was established in an ancestor process that must have
               # subsequently forked. We can't reuse the connection, but we can copy
               # the specification and establish a new connection with it.
-              establish_connection(ancestor_pool.spec).tap do |pool|
+              spec = ancestor_pool.spec
+              spec = spec.to_hash if spec.respond_to?(:to_hash)
+              establish_connection(spec).tap do |pool|
                 pool.schema_cache = ancestor_pool.schema_cache if ancestor_pool.schema_cache
               end
             else

lib/active_record/turntable/active_record_ext/fixtures.rb

@@ -2,6 +2,7 @@
 # force TestFixtures to begin transaction with all shards.
 #
 require "active_record/fixtures"
+require "active_record/turntable/util"
 
 module ActiveRecord
   class FixtureSet
@@ -79,6 +80,7 @@ def setup_fixtures(config = ActiveRecord::Base)
       @fixture_cache = {}
       @fixture_connections = []
       @@already_loaded_fixtures ||= {}
+      @connection_subscriber = nil
 
       # Load fixtures once and begin transaction.
       if run_in_transaction?
@@ -88,11 +90,38 @@ def setup_fixtures(config = ActiveRecord::Base)
           @loaded_fixtures = load_fixtures(config)
           @@already_loaded_fixtures[self.class] = @loaded_fixtures
         end
+
+        # Begin transactions for connections already established
         ActiveRecord::Base.force_connect_all_shards!
         @fixture_connections = enlist_fixture_connections
         @fixture_connections.each do |connection|
           connection.begin_transaction joinable: false
+          if ActiveRecord::Turntable::Util.ar51_or_later?
+            connection.pool.lock_thread = true
+          end
         end
+
+        if ActiveRecord::Turntable::Util.ar51_or_later?
+          # When connections are established in the future, begin a transaction too
+          @connection_subscriber = ActiveSupport::Notifications.subscribe("!connection.active_record") do |_, _, _, _, payload|
+            spec_name = payload[:spec_name] if payload.key?(:spec_name)
+
+            if spec_name
+              begin
+                connection = ActiveRecord::Base.connection_handler.retrieve_connection(spec_name)
+              rescue ConnectionNotEstablished
+                connection = nil
+              end
+
+              if connection && !@fixture_connections.include?(connection)
+                connection.begin_transaction joinable: false
+                connection.pool.lock_thread = true
+                @fixture_connections << connection
+              end
+            end
+          end
+        end
+
       # Load fixtures for every test.
       else
         ActiveRecord::FixtureSet.reset_cache

lib/active_record/turntable/active_record_ext/locking_optimistic.rb

@@ -1,53 +1,105 @@
 module ActiveRecord::Turntable
   module ActiveRecordExt
     module LockingOptimistic
-      ::ActiveRecord::Locking::Optimistic.class_eval <<-EOD
-        private
-        # @note Override to add sharding condition on optimistic locking
-        def _update_record(attribute_names = self.attribute_names) #:nodoc:
-          return super unless locking_enabled?
-          return 0 if attribute_names.empty?
-
-          klass = self.class
-          lock_col = self.class.locking_column
-          previous_lock_value = send(lock_col).to_i
-          increment_lock
-
-          attribute_names += [lock_col]
-          attribute_names.uniq!
-
-          begin
-            relation = self.class.unscoped
-
-            condition_scope = relation.where(
-              self.class.primary_key => id,
-              lock_col => previous_lock_value,
-            )
-            if klass.turntable_enabled? and klass.primary_key != klass.turntable_shard_key.to_s
-              condition_scope = condition_scope.where(
-                klass.turntable_shard_key => self.send(klass.turntable_shard_key)
+      if Util.ar51_or_later?
+        ::ActiveRecord::Locking::Optimistic.class_eval <<-EOD
+          private
+          # @note Override to add sharding condition on optimistic locking
+          def _update_record(attribute_names = self.attribute_names)
+            return super unless locking_enabled?
+            return 0 if attribute_names.empty?
+
+            begin
+              klass = self.class
+
+              lock_col = self.class.locking_column
+
+              previous_lock_value = read_attribute_before_type_cast(lock_col)
+
+              increment_lock
+
+              attribute_names.push(lock_col)
+
+              relation = self.class.unscoped
+
+              condition_scope = relation.where(
+                self.class.primary_key => id,
+                lock_col => previous_lock_value
               )
-            end
+              if klass.turntable_enabled? and klass.primary_key != klass.turntable_shard_key.to_s
+                condition_scope = condition_scope.where(
+                  klass.turntable_shard_key => self.send(klass.turntable_shard_key)
+                )
+              end
+
+              affected_rows = condition_scope.update_all(
+                attributes_for_update(attribute_names).map do |name|
+                  [name, _read_attribute(name)]
+                end.to_h
+              )
+
+              unless affected_rows == 1
+                raise ActiveRecord::StaleObjectError.new(self, "update")
+              end
 
-            affected_rows = condition_scope.update_all(
-              attributes_for_update(attribute_names).map do |name|
-                [name, _read_attribute(name)]
-              end.to_h
-            )
+              affected_rows
 
-            unless affected_rows == 1
-              raise ActiveRecord::StaleObjectError.new(self, "update")
+            # If something went wrong, revert the locking_column value.
+            rescue Exception
+              send(lock_col + "=", previous_lock_value.to_i)
+              raise
             end
+          end
+        EOD
+      else
+        ::ActiveRecord::Locking::Optimistic.class_eval <<-EOD
+          private
+          # @note Override to add sharding condition on optimistic locking
+          def _update_record(attribute_names = self.attribute_names) #:nodoc:
+            return super unless locking_enabled?
+            return 0 if attribute_names.empty?
+
+            klass = self.class
+            lock_col = self.class.locking_column
+            previous_lock_value = send(lock_col).to_i
+            increment_lock
+
+            attribute_names += [lock_col]
+            attribute_names.uniq!
+
+            begin
+              relation = self.class.unscoped
+
+              condition_scope = relation.where(
+                self.class.primary_key => id,
+                lock_col => previous_lock_value,
+              )
+              if klass.turntable_enabled? and klass.primary_key != klass.turntable_shard_key.to_s
+                condition_scope = condition_scope.where(
+                  klass.turntable_shard_key => self.send(klass.turntable_shard_key)
+                )
+              end
+
+              affected_rows = condition_scope.update_all(
+                attributes_for_update(attribute_names).map do |name|
+                  [name, _read_attribute(name)]
+                end.to_h
+              )
 
-            affected_rows
+              unless affected_rows == 1
+                raise ActiveRecord::StaleObjectError.new(self, "update")
+              end
 
-            # If something went wrong, revert the version.
-          rescue Exception
-            send(lock_col + '=', previous_lock_value)
-            raise
+              affected_rows
+
+              # If something went wrong, revert the version.
+            rescue Exception
+              send(lock_col + '=', previous_lock_value)
+              raise
+            end
           end
-        end
-      EOD
+        EOD
+      end
     end
   end
 end