From a86fc93fe1896cf4b32faa54c4e142826a72a109 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Wed, 9 Mar 2022 11:01:43 -0500
Subject: [PATCH] Added select and reselect methods on relation

---
 lib/searchkick/relation.rb | 28 ++++++++++++++++++++++++++++
 test/select_test.rb        | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 0 deletions(-)

diff --git a/lib/searchkick/relation.rb b/lib/searchkick/relation.rb
index a2472f5..287fcd5 100644
--- a/lib/searchkick/relation.rb
+++ b/lib/searchkick/relation.rb
@@ -148,6 +148,34 @@ module Searchkick
     end
 
     # experimental
+    def select(*values, &block)
+      if block_given?
+        private_execute.select(*values, &block)
+      else
+        clone.select!(*values)
+      end
+    end
+
+    # experimental
+    def select!(*values)
+      check_loaded
+      (@options[:select] ||= []).concat(values)
+      self
+    end
+
+    # experimental
+    def reselect(*values)
+      clone.reselect!(*values)
+    end
+
+    # experimental
+    def reselect!(*values)
+      check_loaded
+      @options[:select] = values
+      self
+    end
+
+    # experimental
     def only(*keys)
       Relation.new(@model, @term, **@options.slice(*keys))
     end
diff --git a/test/select_test.rb b/test/select_test.rb
index 5035ac7..c486d95 100644
--- a/test/select_test.rb
+++ b/test/select_test.rb
@@ -9,6 +9,41 @@ class SelectTest < Minitest::Test
     assert_equal 1, result.store_id
   end
 
+  def test_relation
+    store [{name: "Product A", store_id: 1}]
+    result = Product.search("product", load: false).select(:name, :store_id).first
+    assert_equal %w(id name store_id), result.keys.reject { |k| k.start_with?("_") }.sort
+    assert_equal "Product A", result.name
+    assert_equal 1, result.store_id
+  end
+
+  def test_block
+    store [{name: "Product A", store_id: 1}, {name: "Product B", store_id: 2}]
+    assert_equal ["Product B"], Product.search("product", load: false).select { |v| v.store_id == 2 }.map(&:name)
+  end
+
+  def test_block_arguments
+    store [{name: "Product A", store_id: 1}, {name: "Product B", store_id: 2}]
+    error = assert_raises(ArgumentError) do
+      Product.search("product", load: false).select(:name) { |v| v.store_id == 2 }
+    end
+    assert_equal "wrong number of arguments (given 1, expected 0)", error.message
+  end
+  def test_multiple
+    store [{name: "Product A", store_id: 1}]
+    result = Product.search("product", load: false).select(:name).select(:store_id).first
+    assert_equal %w(id name store_id), result.keys.reject { |k| k.start_with?("_") }.sort
+    assert_equal "Product A", result.name
+    assert_equal 1, result.store_id
+  end
+
+  def test_reselect
+    store [{name: "Product A", store_id: 1}]
+    result = Product.search("product", load: false).select(:name).reselect(:store_id).first
+    assert_equal %w(id store_id), result.keys.reject { |k| k.start_with?("_") }.sort
+    assert_equal 1, result.store_id
+  end
+
   def test_array
     store [{name: "Product A", user_ids: [1, 2]}]
     result = Product.search("product", load: false, select: [:user_ids]).first
--
libgit2 0.21.0