From d60a0462d28b0b5a691fe5868be86138566a5fb5 Mon Sep 17 00:00:00 2001
From: Romain Pomier <romain@cardiologs.com>
Date: Wed, 6 Feb 2019 14:29:00 +0100
Subject: [PATCH] Return primary keys when using blocks or values

---
 README.md                       | 25 +++++++++++++++++++------
 lib/bulk_insert.rb              |  6 +++---
 lib/bulk_insert/worker.rb       |  4 ++++
 test/bulk_insert/worker_test.rb |  3 +++
 test/bulk_insert_test.rb        | 15 ++++++++++++++-
 5 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 6b0948b..26b919e 100644
--- a/README.md
+++ b/README.md
@@ -160,6 +160,7 @@ Default value for this option is false.
 destination_columns = [:title, :author]
 
 # Update duplicate rows (MySQL)
+
 Book.bulk_insert(*destination_columns, update_duplicates: true) do |worker|
   worker.add(...)
   worker.add(...)
@@ -175,20 +176,32 @@ end
 
 ### Return Primary Keys (PostgreSQL, PostGIS)
 
-If you want the worker to store primary keys of inserted records, then you can
-use the _return_primary_keys_ option. The worker will store a `result_sets`
-array of `ActiveRecord::Result` objects. Each `ActiveRecord::Result` object
-will contain the primary keys of a batch of inserted records.
+If you want the worker to store primary keys of inserted records, then you
+can use the _return_primary_keys_ option. The `bulk_insert` block will the
+list of primary keys values inserted.
 
 ```ruby
-worker = Book.bulk_insert(*destination_columns, return_primary_keys: true) do
-|worker|
+inserted_ids = Book.bulk_insert(return_primary_keys: true) do |worker|
   worker.add(...)
   worker.add(...)
   # ...
 end
+```
+
+When working with the worker instance directly, the primary keys of the
+inserted records will be stored in the `results_sets` attribute as an array
+of `ActiveRecord::Result` objects. Each `ActiveRecord::Result` object will
+contain the primary keys of a batch of inserted records.
+
+The `inserted_ids` method will return these ids as a flat list.
+
+```ruby
+worker = Book.bulk_insert(*destination_columns, return_primary_keys: true)
+worker.add(...)
+worker.save!
 
 worker.result_sets
+worker.inserted_ids
 ```
 
 ## Ruby and Rails Versions Supported
diff --git a/lib/bulk_insert.rb b/lib/bulk_insert.rb
index 7907aff..068f7f0 100644
--- a/lib/bulk_insert.rb
+++ b/lib/bulk_insert.rb
@@ -8,18 +8,18 @@ def bulk_insert(*columns, values: nil, set_size:500, ignore: false, update_dupli
       columns = default_bulk_columns if columns.empty?
       worker = BulkInsert::Worker.new(connection, table_name, primary_key, columns, set_size, ignore, update_duplicates, return_primary_keys)
 
-      if values.present?
+      if !values.nil?
         transaction do
           worker.add_all(values)
           worker.save!
         end
-        nil
+        return_primary_keys ? worker.inserted_ids : nil
       elsif block_given?
         transaction do
           yield worker
           worker.save!
         end
-        nil
+        return_primary_keys ? worker.inserted_ids : nil
       else
         worker
       end
diff --git a/lib/bulk_insert/worker.rb b/lib/bulk_insert/worker.rb
index 26e0b07..4960f06 100644
--- a/lib/bulk_insert/worker.rb
+++ b/lib/bulk_insert/worker.rb
@@ -36,6 +36,10 @@ def initialize(connection, table_name, primary_key, column_names, set_size=500,
       @set = []
     end
 
+    def inserted_ids
+      @return_primary_keys ? @result_sets.map(&:rows).flatten : nil
+    end
+
     def pending?
       @set.any?
     end
diff --git a/test/bulk_insert/worker_test.rb b/test/bulk_insert/worker_test.rb
index 1ab7c2d..38495e0 100644
--- a/test/bulk_insert/worker_test.rb
+++ b/test/bulk_insert/worker_test.rb
@@ -132,6 +132,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     @insert.save!
 
     assert_equal 0, @insert.result_sets.count
+    assert_nil @insert.inserted_ids
   end
 
 
@@ -159,11 +160,13 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     worker.add greeting: "second"
     worker.save!
     assert_equal 1, worker.result_sets.count
+    assert_equal [], worker.inserted_ids
 
     worker.add greeting: "third"
     worker.add greeting: "fourth"
     worker.save!
     assert_equal 2, worker.result_sets.count
+    assert_equal [], worker.inserted_ids
   end
 
   test "initialized with empty result sets array" do
diff --git a/test/bulk_insert_test.rb b/test/bulk_insert_test.rb
index a31394f..6b1ac3d 100644
--- a/test/bulk_insert_test.rb
+++ b/test/bulk_insert_test.rb
@@ -44,13 +44,26 @@ class BulkInsertTest < ActiveSupport::TestCase
 
   test "bulk_insert with array should save the array immediately" do
     assert_difference "Testing.count", 2 do
-      Testing.bulk_insert values: [
+      result = Testing.bulk_insert values: [
         [ "Hello", 15, true ],
         { greeting: "Hey", age: 20, happy: false }
       ]
+      assert_nil result
     end
   end
 
+  test "bulk_insert with an empty value should not return worker" do
+    assert_nil Testing.bulk_insert values: []
+  end
+
+  test "bulk_insert with option to return primary keys and values should return the new ids" do
+    output = Testing.bulk_insert(
+      values: [["Hello", 15, true]],
+      return_primary_keys: true
+    )
+    assert_equal output, []
+  end
+
   test "default_bulk_columns should return all columns without id" do
     default_columns = %w(greeting age happy created_at updated_at color)