From 50fb2234838d7d455a01b9cf8d0ad270f3512b98 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Mon, 7 Oct 2024 14:26:49 -0700
Subject: [PATCH] Fixed connection leasing for Active Record 7.2+

---
 CHANGELOG.md               | 4 ++++
 lib/rover/data_frame.rb    | 2 +-
 test/active_record_test.rb | 9 +++++++++
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4e0099b..3aa184b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.4.1 (unreleased)
+
+- Fixed connection leasing for Active Record 7.2+
+
 ## 0.4.0 (2024-06-12)
 
 - Added `include?` method to vectors
diff --git a/lib/rover/data_frame.rb b/lib/rover/data_frame.rb
index 6deccd5..7dc48c4 100644
--- a/lib/rover/data_frame.rb
+++ b/lib/rover/data_frame.rb
@@ -41,7 +41,7 @@ def initialize(*args)
           @vectors[k] = to_vector(v, type: types[k])
         end
       elsif defined?(ActiveRecord) && (data.is_a?(ActiveRecord::Relation) || (data.is_a?(Class) && data < ActiveRecord::Base) || data.is_a?(ActiveRecord::Result))
-        result = data.is_a?(ActiveRecord::Result) ? data : data.connection.select_all(data.all.to_sql)
+        result = data.is_a?(ActiveRecord::Result) ? data : data.connection_pool.with_connection { |c| c.select_all(data.all.to_sql) }
         result.columns.each_with_index do |k, i|
           @vectors[k] = to_vector(result.rows.map { |r| r[i] }, type: types[k])
         end
diff --git a/test/active_record_test.rb b/test/active_record_test.rb
index 0f57713..a8146dd 100644
--- a/test/active_record_test.rb
+++ b/test/active_record_test.rb
@@ -28,4 +28,13 @@ def test_result
     assert_vector users.map(&:id), df["id"]
     assert_vector users.map(&:name), df["name"]
   end
+
+  def test_connection_leasing
+    ActiveRecord::Base.connection_handler.clear_active_connections!
+    assert_nil ActiveRecord::Base.connection_pool.active_connection?
+    ActiveRecord::Base.connection_pool.with_connection do
+      Rover::DataFrame.new(User.order(:id))
+    end
+    assert_nil ActiveRecord::Base.connection_pool.active_connection?
+  end
 end