From bcacf9fd5cc2732be29cc07790b256afa79335d0 Mon Sep 17 00:00:00 2001
From: Joshua Young <djry1999@gmail.com>
Date: Fri, 5 Apr 2024 17:39:54 +1000
Subject: [PATCH] Support disallowing permanent connection checkouts with Rails
 7.2

---
 lib/immigrant.rb               | 20 ++++++++++++++++----
 lib/immigrant/key_validator.rb |  6 +++++-
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/lib/immigrant.rb b/lib/immigrant.rb
index bd00647..c72f29b 100644
--- a/lib/immigrant.rb
+++ b/lib/immigrant.rb
@@ -51,13 +51,25 @@ def key_validator
       end
 
       def tables
-        @tables ||= ActiveRecord::Base.connection.tables
+        @tables ||= with_connection(&:tables)
       end
 
       def current_foreign_keys
-        tables.map{ |table|
-          ActiveRecord::Base.connection.foreign_keys(table)
-        }.flatten
+        with_connection do |connection|
+          tables.map{ |table|
+            connection.foreign_keys(table)
+          }.flatten
+        end
+      end
+
+      def with_connection
+        if ActiveRecord::Base.respond_to?(:with_connection)
+          ActiveRecord::Base.with_connection do |connection|
+            yield connection
+          end
+        else
+          yield ActiveRecord::Base.connection
+        end
       end
 
       def model_classes
diff --git a/lib/immigrant/key_validator.rb b/lib/immigrant/key_validator.rb
index 8adad2d..c1a0504 100644
--- a/lib/immigrant/key_validator.rb
+++ b/lib/immigrant/key_validator.rb
@@ -26,7 +26,11 @@ def columns_for(table_name)
     end
 
     def connection
-      @connection ||= ActiveRecord::Base.connection
+      @connection ||= if ActiveRecord::Base.respond_to?(:lease_connection)
+        ActiveRecord::Base.lease_connection
+      else
+        ActiveRecord::Base.connection
+      end
     end
   end
 end