Fixed connection leasing for Active Record 7.2+

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

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

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