From 6c95f8ba77b8db306bde75a0e64f25e6bbdc5e1b Mon Sep 17 00:00:00 2001
From: fumihumi <t.fum1hum@gmail.com>
Date: Tue, 3 Sep 2024 18:08:07 +0900
Subject: [PATCH] wip! support rails 7.2

---
 Rakefile                            |  6 +++++-
 lib/octoball/connection_adapters.rb | 27 +++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/Rakefile b/Rakefile
index 995d4f1..4778582 100644
--- a/Rakefile
+++ b/Rakefile
@@ -43,7 +43,11 @@ namespace :db do
     require './spec/models/application_record'
     ActiveRecord::Base.configurations.configs_for(env_name: "test").each do |config|
       ActiveRecord::Base.establish_connection(config)
-      schema_migration = ActiveRecord::Base.connection.schema_migration
+      schema_migration = if ActiveRecord.gem_version >= Gem::Version.new(7.2)
+                           ActiveRecord::Base.connection.pool.schema_migration
+                         else
+                           ActiveRecord::Base.connection.schema_migration
+                         end
       ActiveRecord::MigrationContext.new("spec/migration", schema_migration)
         .migrate(config.database == 'octoball_shard_5' ? 2 : 1)
     end
diff --git a/lib/octoball/connection_adapters.rb b/lib/octoball/connection_adapters.rb
index 4893129..49c147b 100644
--- a/lib/octoball/connection_adapters.rb
+++ b/lib/octoball/connection_adapters.rb
@@ -13,6 +13,33 @@ module ConnectionHasCurrentShard
     attr_accessor :current_shard
   end
 
+  if ActiveRecord.gem_version >= Gem::Version.new('7.2.0')
+    module ConnectionPoolSetCurrentShard
+      def with_connection(prevent_permanent_checkout: false)
+        lease = connection_lease
+        if lease.connection
+          lease.connection.current_shard = lease.connection.shard
+        end
+
+        super
+      end
+
+      def active_connection?
+        conn = connection_lease.connection
+        conn.current_shard = conn.shard
+        conn
+      end
+
+      def active_connection
+        conn = connection_lease.connection
+        conn.current_shard = conn.shard
+        conn
+      end
+    end
+
+    ::ActiveRecord::ConnectionAdapters::ConnectionPool.prepend(ConnectionPoolSetCurrentShard)
+  end
+
   ::ActiveRecord::ConnectionAdapters::ConnectionHandler.prepend(ConnectionHandlerSetCurrentShard)
   ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ConnectionHasCurrentShard)
 end