Support Rails 6.1

Gemfile

@@ -2,8 +2,8 @@ source "https://rubygems.org"
 
 gemspec
 
-gem "rails", "6.0.0"
+gem "rails", "6.1.0"
 
-gem "actionview", "6.0.0"
-gem "activerecord", "6.0.0"
-gem "activesupport", "6.0.0"
+gem "actionview", "6.1.0"
+gem "activerecord", "6.1.0"
+gem "activesupport", "6.1.0"

lib/active_record/turntable/active_record_ext.rb

@@ -13,6 +13,7 @@ module ActiveRecordExt
       autoload :Sequencer
       autoload :Relation
       autoload :Transactions
+      autoload :AssociationBuilder
       autoload :AssociationPreloader
       autoload :Association
       autoload :LockingOptimistic
@@ -30,6 +31,7 @@ module ActiveRecordExt
       ActiveRecord::Relation.prepend(Relation) unless Util.ar_version_equals_or_later?("5.1.6")
       ActiveRecord::Migration.include(ActiveRecord::Turntable::Migration)
       ActiveRecord::ConnectionAdapters::ConnectionHandler.prepend(ConnectionHandlerExtension)
+      ActiveRecord::Associations::Builder::Association.prepend(AssociationBuilder) if Util.ar61_or_later?
       ActiveRecord::Associations::Preloader::Association.prepend(AssociationPreloader)
       ActiveRecord::Associations::Association.prepend(Association)
       ActiveRecord::QueryCache.prepend(QueryCache)

lib/active_record/turntable/active_record_ext/association.rb

@@ -5,8 +5,10 @@ module ActiveRecordExt
     module Association
       include ShardingCondition
 
-      def self.prepended(mod)
-        ActiveRecord::Associations::Builder::Association::VALID_OPTIONS << :foreign_shard_key
+      unless Util.ar61_or_later?
+        def self.prepended(mod)
+          ActiveRecord::Associations::Builder::Association::VALID_OPTIONS << :foreign_shard_key
+        end
       end
 
       protected

lib/active_record/turntable/active_record_ext/association_builder.rb

@@ -0,0 +1,11 @@
+require "active_record/associations/builder/association"
+
+module ActiveRecord::Turntable
+  module ActiveRecordExt
+    module AssociationBuilder
+      ActiveRecord::Associations::Builder::Association::VALID_OPTIONS = [
+        :class_name, :anonymous_class, :primary_key, :foreign_key, :dependent, :validate, :inverse_of, :strict_loading, :foreign_shard_key
+      ].freeze
+    end
+  end
+end

lib/active_record/turntable/active_record_ext/connection_handler_extension.rb

@@ -6,9 +6,17 @@ def owner_to_turntable_pool
       end
 
       # @note Override not to establish_connection destroy existing connection pool proxy object
-      def retrieve_connection_pool(spec_name)
-        owner_to_turntable_pool.fetch(spec_name) do
-          super
+      if Util.ar61_or_later?
+        def retrieve_connection_pool(owner, role: ActiveRecord::Base.current_role, shard: ActiveRecord::Base.current_shard)
+          owner_to_turntable_pool.fetch(owner) do
+            super
+          end
+        end
+      else
+        def retrieve_connection_pool(spec_name)
+          owner_to_turntable_pool.fetch(spec_name) do
+            super
+          end
         end
       end
     end

lib/active_record/turntable/active_record_ext/database_tasks.rb

@@ -41,7 +41,11 @@ def each_current_turntable_cluster_connected(environment)
           yield(name, configuration)
         end
         ActiveRecord::Base.clear_active_connections!
-        ActiveRecord::Base.establish_connection old_connection_pool.spec.config
+        if ActiveRecord::Turntable::Util.ar61_or_later?
+          ActiveRecord::Base.establish_connection old_connection_pool.db_config
+        else
+          ActiveRecord::Base.establish_connection old_connection_pool.spec.config
+        end
       end
 
       def each_current_turntable_cluster_configuration(environment)
@@ -70,10 +74,21 @@ def each_local_turntable_cluster_configuration
       def current_turntable_cluster_configurations(*environments)
         configurations = []
         environments.each do |environ|
-          config = ActiveRecord::Base.configurations[environ]
+          if ActiveRecord::Turntable::Util.ar61_or_later?
+            config = ActiveRecord::Base.configurations.configs_for(env_name: environ, name: "primary")
+          else
+            config = ActiveRecord::Base.configurations[environ]
+          end
           next unless config
-          %w(shards seq).each do |name|
-            configurations += config[name].to_a if config.has_key?(name)
+
+          if ActiveRecord::Turntable::Util.ar61_or_later?
+            [:shards, :seq].each do |name|
+              configurations += config.configuration_hash[name].to_a if config.configuration_hash.has_key?(name)
+            end
+          else
+            %w(shards seq).each do |name|
+              configurations += config[name].to_a if config.has_key?(name)
+            end
           end
         end
         configurations

lib/active_record/turntable/active_record_ext/fixtures.rb

@@ -83,6 +83,8 @@ def setup_fixtures(config = ActiveRecord::Base)
       @fixture_connections = []
       @@already_loaded_fixtures ||= {}
       @connection_subscriber = nil
+      @legacy_saved_pool_configs = Hash.new { |hash, key| hash[key] = {} }
+      @saved_pool_configs = Hash.new { |hash, key| hash[key] = {} }
 
       # Load fixtures once and begin transaction.
       if run_in_transaction?
@@ -107,18 +109,32 @@ def setup_fixtures(config = ActiveRecord::Base)
           # 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 ActiveRecord::Turntable::Util.ar61_or_later?
+              shard = payload[:shard] if payload.key?(:shard)
+              setup_shared_connection_pool if ActiveRecord::Base.legacy_connection_handling
+            end
 
             if spec_name
               begin
-                connection = ActiveRecord::Base.connection_handler.retrieve_connection(spec_name)
+                if ActiveRecord::Turntable::Util.ar61_or_later?
+                  connection = ActiveRecord::Base.connection_handler.retrieve_connection(spec_name, shard: shard)
+                else
+                  connection = ActiveRecord::Base.connection_handler.retrieve_connection(spec_name)
+                end
               rescue ConnectionNotEstablished
                 connection = nil
               end
 
-              if connection && !@fixture_connections.include?(connection)
-                connection.begin_transaction joinable: false
-                connection.pool.lock_thread = true
-                @fixture_connections << connection
+              if connection
+                if ActiveRecord::Turntable::Util.ar61_or_later?
+                  setup_shared_connection_pool unless ActiveRecord::Base.legacy_connection_handling
+                end
+
+                if !@fixture_connections.include?(connection)
+                  connection.begin_transaction joinable: false
+                  connection.pool.lock_thread = true
+                  @fixture_connections << connection
+                end
               end
             end
           end

lib/active_record/turntable/active_record_ext/transactions.rb

@@ -11,8 +11,16 @@ def with_transaction_returning_status
             self.id.nil? && klass.prefetch_primary_key?
           self.id = klass.next_sequence_value
         end
-        self.class.connection.shards_transaction([self.turntable_shard]) do
-          if Util.ar60_or_later?
+        connection = self.class.connection
+        if Util.ar61_or_later?
+          ensure_finalize = !connection.transaction_open?
+        end
+
+        connection.shards_transaction([self.turntable_shard]) do
+          if Util.ar61_or_later?
+            add_to_transaction(ensure_finalize || has_transactional_callbacks?)
+            remember_transaction_record_state
+          elsif Util.ar60_or_later?
             if has_transactional_callbacks?
               add_to_transaction
             else
@@ -40,19 +48,27 @@ def with_transaction_returning_status
         end
       end
 
-      def add_to_transaction
-        return super unless self.class.turntable_enabled?
+      if Util.ar61_or_later?
+        def add_to_transaction(ensure_finalize = true)
+          return super unless self.class.turntable_enabled?
 
-        if Util.ar60_or_later?
           self.turntable_shard.connection.add_transaction_record(self)
-        else
-          if has_transactional_callbacks?
+        end
+      else
+        def add_to_transaction
+          return super unless self.class.turntable_enabled?
+
+          if Util.ar60_or_later?
             self.turntable_shard.connection.add_transaction_record(self)
           else
-            sync_with_transaction_state
-            set_transaction_state(self.turntable_shard.connection.transaction_state)
+            if has_transactional_callbacks?
+              self.turntable_shard.connection.add_transaction_record(self)
+            else
+              sync_with_transaction_state
+              set_transaction_state(self.turntable_shard.connection.transaction_state)
+            end
+            remember_transaction_record_state
           end
-          remember_transaction_record_state
         end
       end
     end

lib/active_record/turntable/connection_proxy.rb

@@ -280,11 +280,11 @@ def spec
     private
 
       def fixed_shard_entry
-        Thread.current[:turntable_fixed_shard] ||= ThreadSafe::Cache.new
+        Thread.current[:turntable_fixed_shard] ||= Concurrent::Map.new
       end
 
       def current_shard_entry
-        Thread.current[:turntable_current_shard] ||= ThreadSafe::Cache.new
+        Thread.current[:turntable_current_shard] ||= Concurrent::Map.new
       end
   end
 end

lib/active_record/turntable/connection_proxy/mixable.rb

@@ -9,8 +9,8 @@ module Mixable
 
       def mixable?(method, *args)
         (method.to_s =~ METHODS_REGEXP &&
-         args.first !~ EXCLUDE_QUERY_REGEXP) ||
-          (method.to_s == "execute" && args.first =~ QUERY_REGEXP)
+         args.first.to_s !~ EXCLUDE_QUERY_REGEXP) ||
+          (method.to_s == "execute" && args.first.to_s =~ QUERY_REGEXP)
       end
     end
   end

lib/active_record/turntable/pool_proxy.rb

@@ -11,8 +11,14 @@ def with_connection
       yield proxy
     end
 
-    delegate :connected?, :checkout_timeout, :automatic_reconnect, :automatic_reconnect=, :checkout_timeout, :checkout_timeout=, :dead_connection_timeout,
-             :spec, :connections, :size, :reaper, :table_exists?, :query_cache_enabled, :enable_query_cache!, :disable_query_cache!, :schema_cache, :schema_cache=, to: :proxy
+    if Util.ar61_or_later?
+      delegate :connected?, :checkout_timeout, :automatic_reconnect, :automatic_reconnect=, :checkout_timeout, :checkout_timeout=, :dead_connection_timeout,
+      :connections, :size, :reaper, :table_exists?, :query_cache_enabled, :enable_query_cache!, :disable_query_cache!, :schema_cache, :schema_cache=, 
+      :db_config, :pool_config, :connection_klass, :discarded?, :owner_name, to: :proxy
+    else
+      delegate :connected?, :checkout_timeout, :automatic_reconnect, :automatic_reconnect=, :checkout_timeout, :checkout_timeout=, :dead_connection_timeout,
+      :spec, :connections, :size, :reaper, :table_exists?, :query_cache_enabled, :enable_query_cache!, :disable_query_cache!, :schema_cache, :schema_cache=, to: :proxy
+    end
 
     %w(columns_hash column_defaults primary_keys).each do |name|
       define_method(name.to_sym) do

lib/active_record/turntable/seq_shard.rb

@@ -17,7 +17,11 @@ def connection_class_instance
           klass = Class.new(ActiveRecord::Base)
           Connections.const_set(name.classify, klass)
           klass.abstract_class = true
-          klass.establish_connection ActiveRecord::Base.connection_pool.spec.config[:seq][name].with_indifferent_access
+          if Util.ar61_or_later?
+            klass.establish_connection ActiveRecord::Base.connection_pool.db_config.configuration_hash[:seq][name].with_indifferent_access
+          else
+            klass.establish_connection ActiveRecord::Base.connection_pool.spec.config[:seq][name].with_indifferent_access
+          end
         end
         klass
       end

lib/active_record/turntable/shard.rb

@@ -52,7 +52,11 @@ def connection_class_instance
           klass = Class.new(ActiveRecord::Base)
           Connections.const_set(name.classify, klass)
           klass.abstract_class = true
-          klass.establish_connection ActiveRecord::Base.connection_pool.spec.config[:shards][name].with_indifferent_access
+          if Util.ar61_or_later?
+            klass.establish_connection ActiveRecord::Base.connection_pool.db_config.configuration_hash[:shards][name].with_indifferent_access
+          else
+            klass.establish_connection ActiveRecord::Base.connection_pool.spec.config[:shards][name].with_indifferent_access
+          end
         end
         klass
       end

lib/active_record/turntable/sql_tree_patch.rb

@@ -7,6 +7,12 @@ class << self
     attr_accessor :identifier_quote_field_char
   end
   self.identifier_quote_field_char = "`"
+
+  COMMENT_PATTERN = %r{\/\*[\s\S]*?\*\/}.freeze
+  def self.[](query, options = {})
+    sql = query.kind_of?(String) ? query.gsub(COMMENT_PATTERN, "") : query
+    SQLTree::Parser.parse(sql)
+  end
 end
 
 class SQLTree::Token

lib/active_record/turntable/util.rb

@@ -39,6 +39,10 @@ def ar60_or_later?
       ar_version_equals_or_later?("6.0")
     end
 
+    def ar61_or_later?
+      ar_version_equals_or_later?("6.1")
+    end
+
     module_function :ar_version_equals_or_later?,
                     :ar_version_earlier_than?,
                     :ar_version,
@@ -47,6 +51,7 @@ def ar60_or_later?
                     :earlier_than_ar51?,
                     :ar51_or_later?,
                     :ar52_or_later?,
-                    :ar60_or_later?
+                    :ar60_or_later?,
+                    :ar61_or_later?
   end
 end

spec/active_record/turntable/active_record_ext/activerecord_import_ext_spec.rb

@@ -7,8 +7,8 @@
 
     let(:rows) do
       [
-        build(:user_item, user: create(:user, :in_shard1)),
-        build(:user_item, user: create(:user, :in_shard2)),
+        build(:user_item, user: create(:user, :in_shard1), item: create(:item)),
+        build(:user_item, user: create(:user, :in_shard2), item: create(:item)),
       ]
     end
 

spec/active_record/turntable/active_record_ext/locking_optimistic_spec.rb

@@ -11,14 +11,14 @@
   let(:user_profile) { user.user_profile }
 
   describe "optimistic locking" do
-    subject { user_profile.update_attributes(birthday: Date.current) }
+    subject { user_profile.update(birthday: Date.current) }
 
     it { expect { subject }.to change(user_profile, :lock_version).by(1) }
   end
 
   describe "Json serialized column is saved" do
     before do
-      user_profile.update_attributes(data: { foo: "bar" })
+      user_profile.update(data: { foo: "bar" })
       user_profile.reload
     end
 

spec/active_record/turntable/active_record_ext/persistence_spec.rb

@@ -23,7 +23,7 @@
     let(:user) { create(:user, :created_yesterday) }
 
     context "When updating" do
-      subject { user.update_attributes!(nickname: new_nickname) }
+      subject { user.update!(nickname: new_nickname) }
 
       let(:new_nickname) { Faker::Name.unique.name }
 
@@ -73,7 +73,7 @@
     let!(:user_item) { user.user_items.first }
 
     context "When updating" do
-      subject { user_item.update_attributes!(num: 2) }
+      subject { user_item.update!(num: 2) }
 
       it { expect { subject }.not_to raise_error }
 
@@ -124,8 +124,8 @@
   context "When the model is not sharded" do
     let(:item) { create(:item) }
 
-    it { expect { item.update_attributes(name: "hoge") }.not_to raise_error }
-    it { expect { item.update_attributes(name: "hoge") }.to query_like(/WHERE `items`\.`id` = #{item.id}[^\s]*$/) }
+    it { expect { item.update(name: "hoge") }.not_to raise_error }
+    it { expect { item.update(name: "hoge") }.to query_like(/WHERE `items`\.`id` = #{item.id}[^\s]*$/) }
 
     it { expect { item.destroy! }.not_to raise_error }
     it { expect { item.destroy! }.to query_like(/WHERE `items`\.`id` = #{item.id}[^\s]*$/) }

spec/active_record/turntable/active_record_ext/schema_dumper_spec.rb

@@ -10,7 +10,9 @@ def dump_schema
   context "#dump" do
     subject { dump_schema }
 
-    if ActiveRecord::Turntable::Util.ar_version_equals_or_later?("5.0.1")
+    if ActiveRecord::Turntable::Util.ar_version_equals_or_later?("6.1")
+      it { is_expected.to match(/create_sequence_for "users", force: :cascade, charset: "[^"]+", comment: "[^"]+"/) }
+    elsif ActiveRecord::Turntable::Util.ar_version_equals_or_later?("5.0.1")
       it { is_expected.to match(/create_sequence_for "users", force: :cascade, options: "[^"]+", comment: "[^"]+"$/) }
     else
       it { is_expected.to match(/create_sequence_for "users", force: :cascade, options: "[^"]+"/) }