From c870ecb4fd956f43f776705a7a0e833ce5a8d489 Mon Sep 17 00:00:00 2001
From: Michael Newman <5902777+ENewmeration@users.noreply.github.com>
Date: Sun, 3 Nov 2024 13:23:51 -0600
Subject: [PATCH] Allow safe sql values to avoid dangerous query method error
 (#1)

Support plucking custom Arel columns

---------

Co-authored-by: fatkodima <fatkodima123@gmail.com>
---
 CHANGELOG.md                     |  2 ++
 lib/pluck_in_batches/iterator.rb |  9 ++++++++-
 test/pluck_in_batches_test.rb    | 11 +++++++++++
 test/support/schema.rb           |  1 +
 test/test_helper.rb              |  4 ++--
 5 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ac63378..cb21251 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/lib/pluck_in_batches/iterator.rb b/lib/pluck_in_batches/iterator.rb
index 6f0e850..a77fa9d 100644
--- a/lib/pluck_in_batches/iterator.rb
+++ b/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?)
diff --git a/test/pluck_in_batches_test.rb b/test/pluck_in_batches_test.rb
index fb88bde..94416fd 100644
--- a/test/pluck_in_batches_test.rb
+++ b/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
diff --git a/test/support/schema.rb b/test/support/schema.rb
index 442e34b..84895c9 100644
--- a/test/support/schema.rb
+++ b/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|
diff --git a/test/test_helper.rb b/test/test_helper.rb
index c7561ca..fe86acf 100644
--- a/test/test_helper.rb
+++ b/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|