wrap with on_load(:active_record) hook

lib/octoball.rb

@@ -2,29 +2,15 @@
 
 require 'active_record'
 require 'octoball/version'
-require 'octoball/relation_proxy'
-require 'octoball/connection_adapters'
-require 'octoball/connection_handling'
-require 'octoball/current_shard_tracker'
-require 'octoball/association'
-require 'octoball/association_shard_check'
-require 'octoball/persistence'
-require 'octoball/log_subscriber'
 
-class Octoball
-  def self.using(shard, &block)
-    ActiveRecord::Base.connected_to(role: current_role, shard: shard&.to_sym, &block)
-  end
-
-  def self.current_role
-    ActiveRecord::Base.current_role || ActiveRecord::Base.writing_role
-  end
-
-  module UsingShard
-    def using(shard)
-      Octoball::RelationProxy.new(all, shard&.to_sym)
-    end
-  end
-
-  ::ActiveRecord::Base.singleton_class.prepend(UsingShard)
+ActiveSupport.on_load(:active_record) do
+  require 'octoball/relation_proxy'
+  require 'octoball/connection_adapters'
+  require 'octoball/connection_handling'
+  require 'octoball/current_shard_tracker'
+  require 'octoball/association_shard_check'
+  require 'octoball/persistence'
+  require 'octoball/association'
+  require 'octoball/log_subscriber'
+  require 'octoball/using_shard'
 end

lib/octoball/using_shard.rb

@@ -0,0 +1,17 @@
+class Octoball
+  def self.using(shard, &block)
+    ActiveRecord::Base.connected_to(role: current_role, shard: shard&.to_sym, &block)
+  end
+
+  def self.current_role
+    ActiveRecord::Base.current_role || ActiveRecord::Base.writing_role
+  end
+
+  module UsingShard
+    def using(shard)
+      Octoball::RelationProxy.new(all, shard&.to_sym)
+    end
+  end
+
+  ::ActiveRecord::Base.singleton_class.prepend(UsingShard)
+end