Return primary keys when using blocks or values

Author: pomier

README.md

@@ -160,6 +160,7 @@ values. Default value for this option is false.
 destination_columns = [:title, :author]
 
 # Update duplicate rows
+
 Book.bulk_insert(*destination_columns, update_duplicates: true) do |worker|
   worker.add(...)
   worker.add(...)
@@ -169,23 +170,34 @@ 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
 ```
 
-
 ## License
 
 BulkInsert is released under the MIT license (see MIT-LICENSE) by

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

lib/bulk_insert/worker.rb

@@ -32,6 +32,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

test/bulk_insert/worker_test.rb

@@ -128,6 +128,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     @insert.save!
 
     assert_equal 0, @insert.result_sets.count
+    assert_nil @insert.inserted_ids
   end
 
 
@@ -151,11 +152,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

test/bulk_insert_test.rb

@@ -36,13 +36,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)