From ff7bfca07db3fb2529809b5f65f64fa2bcf53d9a Mon Sep 17 00:00:00 2001
From: Peter Loomis <peter.loomis@procore.com>
Date: Thu, 8 Mar 2018 21:40:37 -0800
Subject: [PATCH 01/11] return ids in BulkInsert::Worker

---
 Gemfile                         |  4 ++++
 lib/bulk_insert.rb              |  2 --
 lib/bulk_insert/worker.rb       | 19 +++++++++++++++++--
 test/bulk_insert/worker_test.rb | 25 ++++++++++++++++++++++---
 test/dummy/config/database.yml  | 15 +++------------
 5 files changed, 46 insertions(+), 19 deletions(-)

diff --git a/Gemfile b/Gemfile
index ebfe98b..5cd9b94 100644
--- a/Gemfile
+++ b/Gemfile
@@ -13,3 +13,7 @@ gemspec
 # To use a debugger
 # gem 'byebug', group: [:development, :test]
 
+group :test do
+  gem "pg", "< 1.0"
+end
+
diff --git a/lib/bulk_insert.rb b/lib/bulk_insert.rb
index 6f207b3..88ac868 100644
--- a/lib/bulk_insert.rb
+++ b/lib/bulk_insert.rb
@@ -13,13 +13,11 @@ def bulk_insert(*columns, values: nil, set_size:500, ignore: false, update_dupli
           worker.add_all(values)
           worker.save!
         end
-        nil
       elsif block_given?
         transaction do
           yield worker
           worker.save!
         end
-        nil
       else
         worker
       end
diff --git a/lib/bulk_insert/worker.rb b/lib/bulk_insert/worker.rb
index ba428fc..cd11182 100644
--- a/lib/bulk_insert/worker.rb
+++ b/lib/bulk_insert/worker.rb
@@ -76,12 +76,18 @@ def after_save(&block)
     def save!
       if pending?
         @before_save_callback.(@set) if @before_save_callback
-        compose_insert_query.tap { |query| @connection.execute(query) if query }
+        result = execute_query
         @after_save_callback.() if @after_save_callback
         @set.clear
       end
 
-      self
+      result || []
+    end
+
+    def execute_query
+      if query = compose_insert_query
+        @connection.execute(query)
+      end
     end
 
     def compose_insert_query
@@ -107,6 +113,7 @@ def compose_insert_query
       if !rows.empty?
         sql << rows.join(",")
         sql << on_conflict_statement
+        sql << primary_key_return_statement
         sql
       else
         false
@@ -130,6 +137,14 @@ def insert_ignore
       end
     end
 
+    def primary_key_return_statement
+      if adapter_name =~ /\APostgreSQL/i
+        ' RETURNING id'
+      else
+        ''
+      end
+    end
+
     def on_conflict_statement
       if (adapter_name =~ /\APost(?:greSQL|GIS)/i && ignore )
         ' ON CONFLICT DO NOTHING'
diff --git a/test/bulk_insert/worker_test.rb b/test/bulk_insert/worker_test.rb
index 35622d6..496cdce 100644
--- a/test/bulk_insert/worker_test.rb
+++ b/test/bulk_insert/worker_test.rb
@@ -104,6 +104,10 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     end
   end
 
+  test "save! when not pending should return empty array" do
+    assert_equal [], @insert.save!
+  end
+
   test "save! inserts pending records" do
     @insert.add ["Yo", 15, false, @now, @now]
     @insert.add ["Hello", 25, true, @now, @now]
@@ -121,6 +125,21 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     assert_equal true, hello.happy?
   end
 
+  test "save! returns ids of inserted records" do
+    @insert.add ["Yo"]
+    @insert.add ["Hello"]
+    result_set = @insert.save!
+
+    yo = Testing.find_by(greeting: 'Yo')
+    hello = Testing.find_by(greeting: 'Hello')
+
+    assert_equal result_set.count, 2
+
+    ids_of_inserted_records = result_set.map { |result| result.fetch("id") }
+    assert_includes ids_of_inserted_records, yo.id.to_s
+    assert_includes ids_of_inserted_records, hello.id.to_s
+  end
+
   test "save! calls the after_save handler" do
     x = 41
 
@@ -214,11 +233,11 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
   end
 
   test "adapter dependent default methods" do
-    assert_equal @insert.adapter_name, 'SQLite'
+    assert_equal @insert.adapter_name, 'PostgreSQL'
     assert_equal @insert.insert_sql_statement, "INSERT  INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES "
 
     @insert.add ["Yo", 15, false, nil, nil]
-    assert_equal @insert.compose_insert_query, "INSERT  INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,'f',NULL,NULL,'chartreuse')"
+    assert_equal @insert.compose_insert_query, "INSERT  INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,'f',NULL,NULL,'chartreuse') RETURNING id"
   end
 
   test "adapter dependent mysql methods" do
@@ -284,7 +303,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     pgsql_worker.adapter_name = 'PostgreSQL'
     pgsql_worker.add ["Yo", 15, false, nil, nil]
 
-    assert_equal pgsql_worker.compose_insert_query, "INSERT  INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,'f',NULL,NULL,'chartreuse') ON CONFLICT DO NOTHING"
+    assert_equal pgsql_worker.compose_insert_query, "INSERT  INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,'f',NULL,NULL,'chartreuse') ON CONFLICT DO NOTHING RETURNING id"
   end
 
   test "adapter dependent PostGIS methods" do
diff --git a/test/dummy/config/database.yml b/test/dummy/config/database.yml
index 1c1a37c..cfddb5a 100644
--- a/test/dummy/config/database.yml
+++ b/test/dummy/config/database.yml
@@ -1,25 +1,16 @@
-# SQLite version 3.x
-#   gem install sqlite3
-#
-#   Ensure the SQLite 3 gem is defined in your Gemfile
-#   gem 'sqlite3'
-#
 default: &default
-  adapter: sqlite3
+  adapter: postgresql
   pool: 5
   timeout: 5000
 
 development:
   <<: *default
-  database: db/development.sqlite3
+  database: bulk_insert_development
 
 # Warning: The database defined as "test" will be erased and
 # re-generated from your development database when you run "rake".
 # Do not set this db to the same as development or production.
 test:
   <<: *default
-  database: db/test.sqlite3
+  database: bulk_insert_test
 
-production:
-  <<: *default
-  database: db/production.sqlite3

From 49767415a3dff072c4a1ec3b8de588256df449d0 Mon Sep 17 00:00:00 2001
From: Peter Loomis <peter.loomis@procore.com>
Date: Thu, 8 Mar 2018 22:01:33 -0800
Subject: [PATCH 02/11] integration tests

---
 lib/bulk_insert.rb       |  2 +-
 test/bulk_insert_test.rb | 30 +++++++++++++++++++++++++++---
 test/dummy/db/schema.rb  |  3 +++
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/lib/bulk_insert.rb b/lib/bulk_insert.rb
index 88ac868..8748d60 100644
--- a/lib/bulk_insert.rb
+++ b/lib/bulk_insert.rb
@@ -19,7 +19,7 @@ def bulk_insert(*columns, values: nil, set_size:500, ignore: false, update_dupli
           worker.save!
         end
       else
-        worker
+        []
       end
     end
 
diff --git a/test/bulk_insert_test.rb b/test/bulk_insert_test.rb
index 9b779a6..d5e1a53 100644
--- a/test/bulk_insert_test.rb
+++ b/test/bulk_insert_test.rb
@@ -1,9 +1,9 @@
 require 'test_helper'
 
 class BulkInsertTest < ActiveSupport::TestCase
-  test "bulk_insert without block should return worker" do
+  test "bulk_insert without block should return empty result" do
     result = Testing.bulk_insert
-    assert_kind_of BulkInsert::Worker, result
+    assert_empty result
   end
 
   test "bulk_insert with block should yield worker" do
@@ -23,12 +23,36 @@ 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: [
-        [ "Hello", 15, true, "green" ],
+        [ "Hello", 15, true, Time.now, Time.now, "green" ],
         { greeting: "Hey", age: 20, happy: false }
       ]
     end
   end
 
+  test "ids returned in the same order as the records appear in the insert statement" do
+    attributes_for_insertion = (0..99).map { |i| { age: i } }
+    result_set = Testing.bulk_insert values: attributes_for_insertion
+
+    returned_ids = result_set.map {|result| result.fetch("id").to_i }
+    expected_age_for_id_hash = {}
+    returned_ids.map.with_index do |id, age|
+      expected_age_for_id_hash[id] = age
+    end
+
+    new_saved_records = Testing.find(returned_ids)
+    new_saved_records.each do |record|
+      assert_same(expected_age_for_id_hash[record.id], record.age)
+    end
+  end
+
+  test "returns empty array if called with empty block" do
+    assert_empty(Testing.bulk_insert { |worker| })
+  end
+
+  test "returns empty array if called with empty values" do
+    assert_empty Testing.bulk_insert(values: [])
+  end
+
   test "default_bulk_columns should return all columns without id" do
     default_columns = %w(greeting age happy created_at updated_at color)
 
diff --git a/test/dummy/db/schema.rb b/test/dummy/db/schema.rb
index 5f3a71f..b806206 100644
--- a/test/dummy/db/schema.rb
+++ b/test/dummy/db/schema.rb
@@ -13,6 +13,9 @@
 
 ActiveRecord::Schema.define(version: 20151028194232) do
 
+  # These are extensions that must be enabled in order to support this database
+  enable_extension "plpgsql"
+
   create_table "testings", force: :cascade do |t|
     t.string   "greeting"
     t.integer  "age"

From fef402c877613cdab317ec7f4059e92688bdd599 Mon Sep 17 00:00:00 2001
From: Peter Loomis <peter.loomis@procore.com>
Date: Fri, 9 Mar 2018 15:53:13 -0800
Subject: [PATCH 03/11] undo change to what is returned when empty block

---
 lib/bulk_insert.rb       |  2 +-
 test/bulk_insert_test.rb | 12 ++----------
 2 files changed, 3 insertions(+), 11 deletions(-)

diff --git a/lib/bulk_insert.rb b/lib/bulk_insert.rb
index 8748d60..88ac868 100644
--- a/lib/bulk_insert.rb
+++ b/lib/bulk_insert.rb
@@ -19,7 +19,7 @@ def bulk_insert(*columns, values: nil, set_size:500, ignore: false, update_dupli
           worker.save!
         end
       else
-        []
+        worker
       end
     end
 
diff --git a/test/bulk_insert_test.rb b/test/bulk_insert_test.rb
index d5e1a53..ca58999 100644
--- a/test/bulk_insert_test.rb
+++ b/test/bulk_insert_test.rb
@@ -1,9 +1,9 @@
 require 'test_helper'
 
 class BulkInsertTest < ActiveSupport::TestCase
-  test "bulk_insert without block should return empty result" do
+  test "bulk_insert without block should return worker" do
     result = Testing.bulk_insert
-    assert_empty result
+    assert_kind_of BulkInsert::Worker, result
   end
 
   test "bulk_insert with block should yield worker" do
@@ -45,14 +45,6 @@ class BulkInsertTest < ActiveSupport::TestCase
     end
   end
 
-  test "returns empty array if called with empty block" do
-    assert_empty(Testing.bulk_insert { |worker| })
-  end
-
-  test "returns empty array if called with empty values" do
-    assert_empty Testing.bulk_insert(values: [])
-  end
-
   test "default_bulk_columns should return all columns without id" do
     default_columns = %w(greeting age happy created_at updated_at color)
 

From 8f48764597d24d7f32a921a8b018aeec6465c881 Mon Sep 17 00:00:00 2001
From: Peter Loomis <peter.loomis@procore.com>
Date: Fri, 9 Mar 2018 20:53:32 -0800
Subject: [PATCH 04/11] fix liskov violation

---
 lib/bulk_insert.rb              |  2 ++
 lib/bulk_insert/worker.rb       | 11 +++---
 test/bulk_insert/worker_test.rb | 59 +++++++++++++++++++++++++--------
 test/bulk_insert_test.rb        | 16 ---------
 4 files changed, 54 insertions(+), 34 deletions(-)

diff --git a/lib/bulk_insert.rb b/lib/bulk_insert.rb
index 88ac868..6f207b3 100644
--- a/lib/bulk_insert.rb
+++ b/lib/bulk_insert.rb
@@ -13,11 +13,13 @@ def bulk_insert(*columns, values: nil, set_size:500, ignore: false, update_dupli
           worker.add_all(values)
           worker.save!
         end
+        nil
       elsif block_given?
         transaction do
           yield worker
           worker.save!
         end
+        nil
       else
         worker
       end
diff --git a/lib/bulk_insert/worker.rb b/lib/bulk_insert/worker.rb
index cd11182..c1cba72 100644
--- a/lib/bulk_insert/worker.rb
+++ b/lib/bulk_insert/worker.rb
@@ -5,7 +5,7 @@ class Worker
     attr_accessor :before_save_callback
     attr_accessor :after_save_callback
     attr_accessor :adapter_name
-    attr_reader :ignore, :update_duplicates
+    attr_reader :ignore, :update_duplicates, :results
 
     def initialize(connection, table_name, column_names, set_size=500, ignore=false, update_duplicates=false)
       @connection = connection
@@ -26,6 +26,7 @@ def initialize(connection, table_name, column_names, set_size=500, ignore=false,
       @before_save_callback = nil
       @after_save_callback = nil
 
+      @results = []
       @set = []
     end
 
@@ -76,17 +77,19 @@ def after_save(&block)
     def save!
       if pending?
         @before_save_callback.(@set) if @before_save_callback
-        result = execute_query
+        @results += execute_query
         @after_save_callback.() if @after_save_callback
         @set.clear
       end
 
-      result || []
+      self
     end
 
     def execute_query
       if query = compose_insert_query
-        @connection.execute(query)
+        @connection.execute(query).to_a
+      else
+        []
       end
     end
 
diff --git a/test/bulk_insert/worker_test.rb b/test/bulk_insert/worker_test.rb
index 496cdce..19f069c 100644
--- a/test/bulk_insert/worker_test.rb
+++ b/test/bulk_insert/worker_test.rb
@@ -104,10 +104,6 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     end
   end
 
-  test "save! when not pending should return empty array" do
-    assert_equal [], @insert.save!
-  end
-
   test "save! inserts pending records" do
     @insert.add ["Yo", 15, false, @now, @now]
     @insert.add ["Hello", 25, true, @now, @now]
@@ -125,19 +121,54 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     assert_equal true, hello.happy?
   end
 
-  test "save! returns ids of inserted records" do
-    @insert.add ["Yo"]
-    @insert.add ["Hello"]
-    result_set = @insert.save!
+  test "save! adds to results when there are existing results" do
+    worker = BulkInsert::Worker.new(
+      Testing.connection,
+      Testing.table_name,
+      %w(greeting age happy created_at updated_at color)
+    )
+    worker.add greeting: "first"
+    worker.add geeting: "second"
+    worker.save!
+    assert_equal 2, worker.results.count
 
-    yo = Testing.find_by(greeting: 'Yo')
-    hello = Testing.find_by(greeting: 'Hello')
+    worker.add greeting: "third"
+    worker.add greeting: "fourth"
+    worker.save!
+    assert_equal 4, worker.results.count
+  end
+
+  test "save! does not change worker results if there are no pending rows" do
+    assert_no_difference -> { @insert.results.count } do
+      @insert.save!
+    end
+  end
 
-    assert_equal result_set.count, 2
+  test "results in the same order as the records appear in the insert statement" do
+    attributes_for_insertion = (0..20).map { |i| { age: i } }
+    @insert.add_all attributes_for_insertion
+    result_set = @insert.results
 
-    ids_of_inserted_records = result_set.map { |result| result.fetch("id") }
-    assert_includes ids_of_inserted_records, yo.id.to_s
-    assert_includes ids_of_inserted_records, hello.id.to_s
+    returned_ids = result_set.map {|result| result.fetch("id").to_i }
+    expected_age_for_id_hash = {}
+    returned_ids.map.with_index do |id, index|
+      expected_age_for_id_hash[id] = index
+    end
+
+    new_saved_records = Testing.find(returned_ids)
+    new_saved_records.each do |record|
+      assert_same(expected_age_for_id_hash[record.id], record.age)
+    end
+  end
+
+  test "initialized with empty results" do
+    new_worker = BulkInsert::Worker.new(
+      Testing.connection,
+      Testing.table_name,
+      %w(greeting age happy created_at updated_at color)
+    )
+    assert_instance_of(Array, new_worker.results)
+    assert_empty new_worker.results
   end
 
   test "save! calls the after_save handler" do
diff --git a/test/bulk_insert_test.rb b/test/bulk_insert_test.rb
index ca58999..8f46ac9 100644
--- a/test/bulk_insert_test.rb
+++ b/test/bulk_insert_test.rb
@@ -29,22 +29,6 @@ class BulkInsertTest < ActiveSupport::TestCase
     end
   end
 
-  test "ids returned in the same order as the records appear in the insert statement" do
-    attributes_for_insertion = (0..99).map { |i| { age: i } }
-    result_set = Testing.bulk_insert values: attributes_for_insertion
-
-    returned_ids = result_set.map {|result| result.fetch("id").to_i }
-    expected_age_for_id_hash = {}
-    returned_ids.map.with_index do |id, age|
-      expected_age_for_id_hash[id] = age
-    end
-
-    new_saved_records = Testing.find(returned_ids)
-    new_saved_records.each do |record|
-      assert_same(expected_age_for_id_hash[record.id], record.age)
-    end
-  end
-
   test "default_bulk_columns should return all columns without id" do
     default_columns = %w(greeting age happy created_at updated_at color)
 

From 334174d85220d8b3a59aedefdd62327d406b5fb7 Mon Sep 17 00:00:00 2001
From: Peter Loomis <peter.loomis@procore.com>
Date: Fri, 9 Mar 2018 21:13:04 -0800
Subject: [PATCH 05/11] use array concat so as to not allocate a new array on
 every save

---
 lib/bulk_insert/worker.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/bulk_insert/worker.rb b/lib/bulk_insert/worker.rb
index c1cba72..5fe5b79 100644
--- a/lib/bulk_insert/worker.rb
+++ b/lib/bulk_insert/worker.rb
@@ -77,7 +77,7 @@ def after_save(&block)
     def save!
       if pending?
         @before_save_callback.(@set) if @before_save_callback
-        @results += execute_query
+        @results.concat(execute_query)
         @after_save_callback.() if @after_save_callback
         @set.clear
       end

From ad6f4e66ad23e4e4670d747d633c11e2dc8aff03 Mon Sep 17 00:00:00 2001
From: Peter Loomis <peter.loomis@procore.com>
Date: Sat, 10 Mar 2018 12:34:01 -0800
Subject: [PATCH 06/11] store result_sets, not results

---
 lib/bulk_insert/worker.rb       | 11 +++++------
 test/bulk_insert/worker_test.rb | 18 ++++++++++--------
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/lib/bulk_insert/worker.rb b/lib/bulk_insert/worker.rb
index 5fe5b79..7b249f4 100644
--- a/lib/bulk_insert/worker.rb
+++ b/lib/bulk_insert/worker.rb
@@ -5,7 +5,7 @@ class Worker
     attr_accessor :before_save_callback
     attr_accessor :after_save_callback
     attr_accessor :adapter_name
-    attr_reader :ignore, :update_duplicates, :results
+    attr_reader :ignore, :update_duplicates, :result_sets
 
     def initialize(connection, table_name, column_names, set_size=500, ignore=false, update_duplicates=false)
       @connection = connection
@@ -26,7 +26,7 @@ def initialize(connection, table_name, column_names, set_size=500, ignore=false,
       @before_save_callback = nil
       @after_save_callback = nil
 
-      @results = []
+      @result_sets = []
       @set = []
     end
 
@@ -77,7 +77,7 @@ def after_save(&block)
     def save!
       if pending?
         @before_save_callback.(@set) if @before_save_callback
-        @results.concat(execute_query)
+        execute_query
         @after_save_callback.() if @after_save_callback
         @set.clear
       end
@@ -87,9 +87,8 @@ def save!
 
     def execute_query
       if query = compose_insert_query
-        @connection.execute(query).to_a
-      else
-        []
+        result_set = @connection.execute(query)
+        @result_sets.push(result_set)
       end
     end
 
diff --git a/test/bulk_insert/worker_test.rb b/test/bulk_insert/worker_test.rb
index 19f069c..64f0bc9 100644
--- a/test/bulk_insert/worker_test.rb
+++ b/test/bulk_insert/worker_test.rb
@@ -130,16 +130,18 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     worker.add greeting: "first"
     worker.add geeting: "second"
     worker.save!
-    assert_equal 2, worker.results.count
+    assert_equal 1, worker.result_sets.count
+    assert_equal 2, worker.result_sets.map(&:to_a).flatten.count
 
     worker.add greeting: "third"
     worker.add greeting: "fourth"
     worker.save!
-    assert_equal 4, worker.results.count
+    assert_equal 2, worker.result_sets.count
+    assert_equal 4, worker.result_sets.map(&:to_a).flatten.count
   end
 
   test "save! does not change worker results if there are no pending rows" do
-    assert_no_difference -> { @insert.results.count } do
+    assert_no_difference -> { @insert.result_sets.count } do
       @insert.save!
     end
   end
@@ -147,9 +149,9 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
   test "results in the same order as the records appear in the insert statement" do
     attributes_for_insertion = (0..20).map { |i| { age: i } }
     @insert.add_all attributes_for_insertion
-    result_set = @insert.results
+    results = @insert.result_sets.map(&:to_a).flatten
 
-    returned_ids = result_set.map {|result| result.fetch("id").to_i }
+    returned_ids = results.map {|result| result.fetch("id").to_i }
     expected_age_for_id_hash = {}
     returned_ids.map.with_index do |id, index|
       expected_age_for_id_hash[id] = index
@@ -161,14 +163,14 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     end
   end
 
-  test "initialized with empty results" do
+  test "initialized with empty result_sets array" do
     new_worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
       %w(greeting age happy created_at updated_at color)
     )
-    assert_instance_of(Array, new_worker.results)
-    assert_empty new_worker.results
+    assert_instance_of(Array, new_worker.result_sets)
+    assert_empty new_worker.result_sets
   end
 
   test "save! calls the after_save handler" do

From 6eae2c17a78e653e6c76a5e5d314a994915e0508 Mon Sep 17 00:00:00 2001
From: Peter Loomis <peter.loomis@procore.com>
Date: Tue, 13 Mar 2018 15:29:28 -0700
Subject: [PATCH 07/11] fix type in spec

---
 test/bulk_insert/worker_test.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/bulk_insert/worker_test.rb b/test/bulk_insert/worker_test.rb
index 64f0bc9..22bcdfb 100644
--- a/test/bulk_insert/worker_test.rb
+++ b/test/bulk_insert/worker_test.rb
@@ -128,7 +128,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
       %w(greeting age happy created_at updated_at color)
     )
     worker.add greeting: "first"
-    worker.add geeting: "second"
+    worker.add greeting: "second"
     worker.save!
     assert_equal 1, worker.result_sets.count
     assert_equal 2, worker.result_sets.map(&:to_a).flatten.count

From 7c0a06c45b9ef772836fe7616aed4012f31899fe Mon Sep 17 00:00:00 2001
From: Peter Loomis <peter.loomis@procore.com>
Date: Wed, 14 Mar 2018 11:08:15 -0700
Subject: [PATCH 08/11] make return of primary keys opt-in

---
 lib/bulk_insert/worker.rb       |  9 +++--
 test/bulk_insert/worker_test.rb | 68 +++++++++++++++++++++++++++------
 2 files changed, 62 insertions(+), 15 deletions(-)

diff --git a/lib/bulk_insert/worker.rb b/lib/bulk_insert/worker.rb
index 7b249f4..aab04bd 100644
--- a/lib/bulk_insert/worker.rb
+++ b/lib/bulk_insert/worker.rb
@@ -7,7 +7,7 @@ class Worker
     attr_accessor :adapter_name
     attr_reader :ignore, :update_duplicates, :result_sets
 
-    def initialize(connection, table_name, column_names, set_size=500, ignore=false, update_duplicates=false)
+    def initialize(connection, table_name, column_names, set_size=500, ignore=false, update_duplicates=false, return_primary_keys=false)
       @connection = connection
       @set_size = set_size
 
@@ -15,6 +15,7 @@ def initialize(connection, table_name, column_names, set_size=500, ignore=false,
       # INSERT IGNORE only fails inserts with duplicate keys or unallowed nulls not the whole set of inserts
       @ignore = ignore
       @update_duplicates = update_duplicates
+      @return_primary_keys = return_primary_keys
 
       columns = connection.columns(table_name)
       column_map = columns.inject({}) { |h, c| h.update(c.name => c) }
@@ -88,7 +89,7 @@ def save!
     def execute_query
       if query = compose_insert_query
         result_set = @connection.execute(query)
-        @result_sets.push(result_set)
+        @result_sets.push(result_set) if @return_primary_keys
       end
     end
 
@@ -115,7 +116,7 @@ def compose_insert_query
       if !rows.empty?
         sql << rows.join(",")
         sql << on_conflict_statement
-        sql << primary_key_return_statement
+        sql << primary_key_return_statement if @return_primary_keys
         sql
       else
         false
@@ -140,7 +141,7 @@ def insert_ignore
     end
 
     def primary_key_return_statement
-      if adapter_name =~ /\APostgreSQL/i
+      if adapter_name =~ /\APost(?:greSQL|GIS)/i
         ' RETURNING id'
       else
         ''
diff --git a/test/bulk_insert/worker_test.rb b/test/bulk_insert/worker_test.rb
index 22bcdfb..b71f3a2 100644
--- a/test/bulk_insert/worker_test.rb
+++ b/test/bulk_insert/worker_test.rb
@@ -121,11 +121,32 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     assert_equal true, hello.happy?
   end
 
-  test "save! adds to results when there are existing results" do
+  test "save! does not add to result sets when not returning primary keys" do
     worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
-      %w(greeting age happy created_at updated_at color)
+      %w(greeting age happy created_at updated_at color),
+      500,
+      false,
+      false,
+      false
+    )
+    worker.add greeting: "first"
+    worker.add greeting: "second"
+    worker.save!
+
+    assert_equal 0, worker.result_sets.count
+  end
+
+  test "save! adds to result sets when returning primary keys" do
+    worker = BulkInsert::Worker.new(
+      Testing.connection,
+      Testing.table_name,
+      %w(greeting age happy created_at updated_at color),
+      500,
+      false,
+      false,
+      true
     )
     worker.add greeting: "first"
     worker.add greeting: "second"
@@ -140,16 +161,35 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     assert_equal 4, worker.result_sets.map(&:to_a).flatten.count
   end
 
-  test "save! does not change worker results if there are no pending rows" do
-    assert_no_difference -> { @insert.result_sets.count } do
-      @insert.save!
+  test "save! does not change worker result sets if there are no pending rows" do
+    worker = BulkInsert::Worker.new(
+      Testing.connection,
+      Testing.table_name,
+      %w(greeting age happy created_at updated_at color),
+      500,
+      false,
+      false,
+      true
+    )
+    assert_no_difference -> { worker.result_sets.count } do
+      worker.save!
     end
   end
 
   test "results in the same order as the records appear in the insert statement" do
+    worker = BulkInsert::Worker.new(
+      Testing.connection,
+      Testing.table_name,
+      %w(greeting age happy created_at updated_at color),
+      500,
+      false,
+      false,
+      true
+    )
+
     attributes_for_insertion = (0..20).map { |i| { age: i } }
-    @insert.add_all attributes_for_insertion
-    results = @insert.result_sets.map(&:to_a).flatten
+    worker.add_all attributes_for_insertion
+    results = worker.result_sets.map(&:to_a).flatten
 
     returned_ids = results.map {|result| result.fetch("id").to_i }
     expected_age_for_id_hash = {}
@@ -270,7 +310,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     assert_equal @insert.insert_sql_statement, "INSERT  INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES "
 
     @insert.add ["Yo", 15, false, nil, nil]
-    assert_equal @insert.compose_insert_query, "INSERT  INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,'f',NULL,NULL,'chartreuse') RETURNING id"
+    assert_equal @insert.compose_insert_query, "INSERT  INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,'f',NULL,NULL,'chartreuse')"
   end
 
   test "adapter dependent mysql methods" do
@@ -332,7 +372,10 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
       Testing.table_name,
       %w(greeting age happy created_at updated_at color),
       500, # batch size
-      true) # ignore
+      true, # ignore
+      false, # update duplicates
+      true # return primary key
+    )
     pgsql_worker.adapter_name = 'PostgreSQL'
     pgsql_worker.add ["Yo", 15, false, nil, nil]
 
@@ -345,11 +388,14 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
       Testing.table_name,
       %w(greeting age happy created_at updated_at color),
       500, # batch size
-      true) # ignore
+      true, # ignore
+      false, # update duplicates
+      true # return primary key
+    ) # ignore
     pgsql_worker.adapter_name = 'PostGIS'
     pgsql_worker.add ["Yo", 15, false, nil, nil]
 
-    assert_equal pgsql_worker.compose_insert_query, "INSERT  INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,'f',NULL,NULL,'chartreuse') ON CONFLICT DO NOTHING"
+    assert_equal pgsql_worker.compose_insert_query, "INSERT  INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,'f',NULL,NULL,'chartreuse') ON CONFLICT DO NOTHING RETURNING id"
   end
 
   test "adapter dependent sqlite3 methods (with lowercase adapter name)" do

From c4bfb353389af175e7cd57f3463ba8aa981ea411 Mon Sep 17 00:00:00 2001
From: Peter Loomis <peter.loomis@procore.com>
Date: Wed, 14 Mar 2018 11:32:29 -0700
Subject: [PATCH 09/11] make return of primary keys opt-in

---
 lib/bulk_insert.rb              |  4 ++--
 lib/bulk_insert/worker.rb       |  9 +++++----
 test/bulk_insert/worker_test.rb | 14 ++++++++++++++
 test/bulk_insert_test.rb        | 14 ++++++++++++++
 4 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/lib/bulk_insert.rb b/lib/bulk_insert.rb
index 6f207b3..851e7f9 100644
--- a/lib/bulk_insert.rb
+++ b/lib/bulk_insert.rb
@@ -4,9 +4,9 @@ module BulkInsert
   extend ActiveSupport::Concern
 
   module ClassMethods
-    def bulk_insert(*columns, values: nil, set_size:500, ignore: false, update_duplicates: false)
+    def bulk_insert(*columns, values: nil, set_size:500, ignore: false, update_duplicates: false, return_primary_keys: false)
       columns = default_bulk_columns if columns.empty?
-      worker = BulkInsert::Worker.new(connection, table_name, columns, set_size, ignore, update_duplicates)
+      worker = BulkInsert::Worker.new(connection, table_name, self.primary_key, columns, set_size, ignore, update_duplicates, return_primary_keys)
 
       if values.present?
         transaction do
diff --git a/lib/bulk_insert/worker.rb b/lib/bulk_insert/worker.rb
index aab04bd..2c55721 100644
--- a/lib/bulk_insert/worker.rb
+++ b/lib/bulk_insert/worker.rb
@@ -7,7 +7,7 @@ class Worker
     attr_accessor :adapter_name
     attr_reader :ignore, :update_duplicates, :result_sets
 
-    def initialize(connection, table_name, column_names, set_size=500, ignore=false, update_duplicates=false, return_primary_keys=false)
+    def initialize(connection, table_name, primary_key, column_names, set_size=500, ignore=false, update_duplicates=false, return_primary_keys=false)
       @connection = connection
       @set_size = set_size
 
@@ -20,6 +20,7 @@ def initialize(connection, table_name, column_names, set_size=500, ignore=false,
       columns = connection.columns(table_name)
       column_map = columns.inject({}) { |h, c| h.update(c.name => c) }
 
+      @primary_key = primary_key
       @columns = column_names.map { |name| column_map[name.to_s] }
       @table_name = connection.quote_table_name(table_name)
       @column_names = column_names.map { |name| connection.quote_column_name(name) }.join(",")
@@ -116,7 +117,7 @@ def compose_insert_query
       if !rows.empty?
         sql << rows.join(",")
         sql << on_conflict_statement
-        sql << primary_key_return_statement if @return_primary_keys
+        sql << primary_key_return_statement
         sql
       else
         false
@@ -141,8 +142,8 @@ def insert_ignore
     end
 
     def primary_key_return_statement
-      if adapter_name =~ /\APost(?:greSQL|GIS)/i
-        ' RETURNING id'
+      if @return_primary_keys && adapter_name =~ /\APost(?:greSQL|GIS)/i
+        " RETURNING #{@primary_key}"
       else
         ''
       end
diff --git a/test/bulk_insert/worker_test.rb b/test/bulk_insert/worker_test.rb
index b71f3a2..bdd2dff 100644
--- a/test/bulk_insert/worker_test.rb
+++ b/test/bulk_insert/worker_test.rb
@@ -5,6 +5,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     @insert = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color))
     @now = Time.now
   end
@@ -125,6 +126,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color),
       500,
       false,
@@ -142,6 +144,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color),
       500,
       false,
@@ -165,6 +168,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color),
       500,
       false,
@@ -180,6 +184,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color),
       500,
       false,
@@ -207,6 +212,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     new_worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color)
     )
     assert_instance_of(Array, new_worker.result_sets)
@@ -317,6 +323,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     mysql_worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color),
       500, # batch size
       true) # ignore
@@ -336,6 +343,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     mysql_worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color),
       500, # batch size
       true, # ignore
@@ -354,6 +362,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     mysql_worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color),
       500, # batch size
       true) # ignore
@@ -370,6 +379,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     pgsql_worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color),
       500, # batch size
       true, # ignore
@@ -386,6 +396,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     pgsql_worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color),
       500, # batch size
       true, # ignore
@@ -402,6 +413,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     sqlite_worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color),
       500, # batch size
       true) # ignore
@@ -415,6 +427,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     sqlite_worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color),
       500, # batch size
       true) # ignore
@@ -428,6 +441,7 @@ class BulkInsertWorkerTest < ActiveSupport::TestCase
     mysql_worker = BulkInsert::Worker.new(
       Testing.connection,
       Testing.table_name,
+      'id',
       %w(greeting age happy created_at updated_at color),
       500, # batch size
       false, # ignore
diff --git a/test/bulk_insert_test.rb b/test/bulk_insert_test.rb
index 8f46ac9..f033ebe 100644
--- a/test/bulk_insert_test.rb
+++ b/test/bulk_insert_test.rb
@@ -20,6 +20,20 @@ class BulkInsertTest < ActiveSupport::TestCase
     end
   end
 
+  test "worker should not have any result sets without option for returning primary keys" do
+    worker = Testing.bulk_insert
+    worker.add greeting: "hello"
+    worker.save!
+    assert_empty worker.result_sets
+  end
+
+  test "with option to return primary keys, worker should have result sets" do
+    worker = Testing.bulk_insert(return_primary_keys: true)
+    worker.add greeting: "yo"
+    worker.save!
+    assert_equal 1, worker.result_sets.count
+  end
+
   test "bulk_insert with array should save the array immediately" do
     assert_difference "Testing.count", 2 do
       Testing.bulk_insert values: [

From ec27cac760146be86cce694d602f4f7db6bf6a52 Mon Sep 17 00:00:00 2001
From: Peter Loomis <peter.loomis@procore.com>
Date: Wed, 14 Mar 2018 11:34:58 -0700
Subject: [PATCH 10/11] use exec_query instead of execute

---
 lib/bulk_insert/worker.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/bulk_insert/worker.rb b/lib/bulk_insert/worker.rb
index 2c55721..f4ba7c2 100644
--- a/lib/bulk_insert/worker.rb
+++ b/lib/bulk_insert/worker.rb
@@ -89,7 +89,7 @@ def save!
 
     def execute_query
       if query = compose_insert_query
-        result_set = @connection.execute(query)
+        result_set = @connection.exec_query(query)
         @result_sets.push(result_set) if @return_primary_keys
       end
     end

From b355108a49ece9bec379cf0a48bd42be7f3ebcac Mon Sep 17 00:00:00 2001
From: Peter Loomis <peter.loomis@procore.com>
Date: Wed, 14 Mar 2018 13:16:42 -0700
Subject: [PATCH 11/11] remove redundant self

---
 lib/bulk_insert.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/bulk_insert.rb b/lib/bulk_insert.rb
index 851e7f9..7907aff 100644
--- a/lib/bulk_insert.rb
+++ b/lib/bulk_insert.rb
@@ -6,7 +6,7 @@ module BulkInsert
   module ClassMethods
     def bulk_insert(*columns, values: nil, set_size:500, ignore: false, update_duplicates: false, return_primary_keys: false)
       columns = default_bulk_columns if columns.empty?
-      worker = BulkInsert::Worker.new(connection, table_name, self.primary_key, columns, set_size, ignore, update_duplicates, return_primary_keys)
+      worker = BulkInsert::Worker.new(connection, table_name, primary_key, columns, set_size, ignore, update_duplicates, return_primary_keys)
 
       if values.present?
         transaction do