Allow safe sql values to avoid dangerous query method error (#1)

Support plucking custom Arel columns

---------

Co-authored-by: fatkodima <[email protected]>

CHANGELOG.md

@@ -1,5 +1,7 @@
 ## master (unreleased)
 
+- Support plucking custom Arel columns
+
 ## 0.2.0 (2023-07-24)
 
 - Support specifying per cursor column ordering when batching

lib/pluck_in_batches/iterator.rb

@@ -35,7 +35,14 @@ def each_batch(*columns, start: nil, finish: nil, batch_size: 1000, error_on_ign
         raise ArgumentError, ":order must be :asc or :desc or an array consisting of :asc or :desc, got #{order.inspect}"
       end
 
-      pluck_columns = columns.map(&:to_s)
+      pluck_columns = columns.map do |column|
+        if Arel.arel_node?(column)
+          column
+        else
+          column.to_s
+        end
+      end
+
       cursor_columns = Array(cursor_column).map(&:to_s)
       cursor_column_indexes = cursor_column_indexes(pluck_columns, cursor_columns)
       missing_cursor_columns = cursor_column_indexes.count(&:nil?)

test/pluck_in_batches_test.rb

@@ -48,6 +48,17 @@ def test_pluck_in_batches_should_return_batches
     end
   end
 
+  def test_pluck_in_batches_should_return_batches_for_arel_columns
+    ids_and_ranks = User.order(:id).pluck(:id, Arel.sql("json_extract(users.metadata, '$.rank')"))
+
+    assert_queries(User.count + 1) do
+      User.pluck_in_batches(:id, Arel.sql("json_extract(users.metadata, '$.rank')"), batch_size: 1).with_index do |batch, index|
+        assert_kind_of Array, batch
+        assert_equal ids_and_ranks[index], batch.first
+      end
+    end
+  end
+
   def pluck_in_batches_desc_order
     ids = User.order(id: :desc).ids
     assert_queries(User.count + 1) do

test/support/schema.rb

@@ -3,6 +3,7 @@
 ActiveRecord::Schema.define do
   create_table :users, force: true do |t|
     t.string :name
+    t.json :metadata, null: false, default: {}
   end
 
   create_table :subscribers, id: false, primary_key: :nick, force: true do |t|

test/test_helper.rb

@@ -23,9 +23,9 @@ def prepare_database
   # Create users
   values = 20.times.map do |i|
     id = i + 1
-    "(#{id}, 'User-#{id}')"
+    "(#{id}, 'User-#{id}', json('{\"rank\": #{i}}'))"
   end.join(", ")
-  ActiveRecord::Base.connection.execute("INSERT INTO users (id, name) VALUES #{values}")
+  ActiveRecord::Base.connection.execute("INSERT INTO users (id, name, metadata) VALUES #{values}")
 
   # Create subscribers
   values = 10.times.map do |i|