lib/octoshark/connection_pools_manager.rb in octoshark-0.3.0 vs lib/octoshark/connection_pools_manager.rb in octoshark-0.4.0

@@ -69,29 +69,60 @@
         @connection_pools[name] = create_connection_pool(name, config)
       end
     end
 
     def create_connection_pool(name, config)
-      spec =
-        if defined?(ActiveRecord::ConnectionAdapters::PoolConfig)
-          env_name = defined?(Rails) ? Rails.env : nil
-          db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(env_name, name, config)
-          ActiveRecord::ConnectionAdapters::PoolConfig.new(owner_name = ActiveRecord::Base, db_config)
+      spec = build_connection_pool_spec(name, config)
+
+      ActiveRecord::ConnectionAdapters::ConnectionPool.new(spec)
+    end
+
+    private
+
+    def build_connection_pool_spec(name, config)
+      if active_record_6_1_or_7?
+        env_name = defined?(Rails) ? Rails.env : nil
+        db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(env_name, name, config)
+
+        pool_config_class = ActiveRecord::ConnectionAdapters::PoolConfig
+
+        if pool_config_class.instance_method(:initialize).arity == 2
+          # ActiveRecord 6.1
+          pool_config_class.new(owner_name = ActiveRecord::Base, db_config)
         else
-          adapter_method = "#{config[:adapter]}_connection"
-          if defined?(ActiveRecord::ConnectionAdapters::ConnectionSpecification)
-            spec_class = ActiveRecord::ConnectionAdapters::ConnectionSpecification
+          # ActiveRecord 7.0
+          pool_config_class.new(
+            owner_name = ActiveRecord::Base,
+            db_config,
+            role = ActiveRecord::Base.current_role,
+            shard = ActiveRecord::Base.current_shard
+          )
+        end
+      else
+        adapter_method = "#{config[:adapter]}_connection"
 
-            if spec_class.instance_method(:initialize).arity == 3
-              spec_class.new(name, config, adapter_method)
-            else
-              spec_class.new(config, adapter_method)
-            end
+        if active_record_4_or_5_or_6?
+          spec_class = ActiveRecord::ConnectionAdapters::ConnectionSpecification
+
+          if spec_class.instance_method(:initialize).arity == 3
+            # ActiveRecord 5.x and 6.0
+            spec_class.new(name, config, adapter_method)
           else
-            ActiveRecord::Base::ConnectionSpecification.new(config, adapter_method)
+            # ActiveRecord 4.x
+            spec_class.new(config, adapter_method)
           end
+        else
+          # ActiveRecord 3.x
+          ActiveRecord::Base::ConnectionSpecification.new(config, adapter_method)
         end
+      end
+    end
 
-      ActiveRecord::ConnectionAdapters::ConnectionPool.new(spec)
+    def active_record_6_1_or_7?
+      defined?(ActiveRecord::ConnectionAdapters::PoolConfig)
+    end
+
+    def active_record_4_or_5_or_6?
+      defined?(ActiveRecord::ConnectionAdapters::ConnectionSpecification)
     end
   end
 end