Fixed connection leasing for Active Record 7.2 - fixes #26

ankane · Oct 7, 2024 · 29684d8 · 29684d8
1 parent 5f88029
commit 29684d8
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,7 @@
 - Added experimental support for MariaDB 11.6 Vector
 - Added experimental support for MySQL 9
 - Changed `normalize` option to use Active Record normalization
+- Fixed connection leasing for Active Record 7.2
 - Dropped support for Active Record < 7
 
 ## 0.4.3 (2024-09-02)

diff --git a/lib/neighbor/model.rb b/lib/neighbor/model.rb
@@ -88,7 +88,7 @@ def self.neighbor_attributes
 
           distance = distance.to_s
 
-          quoted_attribute = "#{connection.quote_table_name(table_name)}.#{connection.quote_column_name(attribute_name)}"
+          quoted_attribute = with_connection { |c| "#{c.quote_table_name(table_name)}.#{c.quote_column_name(attribute_name)}" }
 
           column_info = columns_hash[attribute_name.to_s]
           column_type = column_info&.type
@@ -113,7 +113,7 @@ def self.neighbor_attributes
           Neighbor::Utils.validate(vector, dimensions: dimensions, type: type || Utils.type(adapter, column_info&.type), adapter: adapter)
           vector = Neighbor::Utils.normalize(vector, column_info: column_info) if normalize
 
-          query = connection.quote(column_attribute.serialize(vector))
+          query = with_connection { |c| c.quote(column_attribute.serialize(vector)) }
 
           if !precision.nil?
             if adapter != :postgresql || column_type != :vector
@@ -124,7 +124,7 @@ def self.neighbor_attributes
             when "half"
               cast_dimensions = dimensions || column_info&.limit
               raise ArgumentError, "Unknown dimensions" unless cast_dimensions
-              quoted_attribute += "::halfvec(#{connection.quote(cast_dimensions.to_i)})"
+              quoted_attribute += "::halfvec(#{with_connection { |c| c.quote(cast_dimensions.to_i) }})"
             else
               raise ArgumentError, "Invalid precision"
             end

diff --git a/test/neighbor_test.rb b/test/neighbor_test.rb
@@ -16,6 +16,15 @@ def test_schema
     assert_match %{t.sparsevec "sparse_embedding", limit: 3}, contents
   end
 
+  def test_connection_pool
+    PostgresRecord.connection_handler.clear_active_connections!
+    assert_nil PostgresRecord.connection_pool.active_connection?
+    PostgresRecord.connection_pool.with_connection do
+      Item.nearest_neighbors(:embedding, [1, 1, 1], distance: "euclidean")
+    end
+    assert_nil PostgresRecord.connection_pool.active_connection?
+  end
+
   def test_composite_primary_key
     skip if ActiveRecord::VERSION::STRING.to_f < 7.1