bulk_insert2 1.0.0

data/lib/bulk_insert/worker.rb

@@ -0,0 +1,136 @@
+require_relative 'statement_adapters'
+
+module BulkInsert
+  class Worker
+    attr_reader :connection
+    attr_accessor :set_size
+    attr_accessor :before_save_callback
+    attr_accessor :after_save_callback
+    attr_accessor :adapter_name
+    attr_reader :ignore, :update_duplicates, :result_sets
+
+    def initialize(connection, table_name, primary_key, column_names, set_size=500, ignore=false, update_duplicates=false, return_primary_keys=false)
+      @statement_adapter = StatementAdapters.adapter_for(connection)
+
+      @connection = connection
+      @set_size = set_size
+
+      @adapter_name = connection.adapter_name
+      # 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) }
+
+      @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(",")
+
+      @before_save_callback = nil
+      @after_save_callback = nil
+
+      @result_sets = []
+      @set = []
+    end
+
+    def pending?
+      @set.any?
+    end
+
+    def pending_count
+      @set.count
+    end
+
+    def add(values)
+      save! if @set.length >= set_size
+
+      values = values.with_indifferent_access if values.is_a?(Hash)
+      mapped = @columns.map.with_index do |column, index|
+          value_exists = values.is_a?(Hash) ? values.key?(column.name) : (index < values.length)
+          if !value_exists
+            if column.default.present?
+              column.default
+            elsif column.name == "created_at" || column.name == "updated_at"
+              :__timestamp_placeholder
+            else
+              nil
+            end
+          else
+            values.is_a?(Hash) ? values[column.name] : values[index]
+          end
+        end
+
+      @set.push(mapped)
+      self
+    end
+
+    def add_all(rows)
+      rows.each { |row| add(row) }
+      self
+    end
+
+    def before_save(&block)
+      @before_save_callback = block
+    end
+
+    def after_save(&block)
+      @after_save_callback = block
+    end
+
+    def save!
+      if pending?
+        @before_save_callback.(@set) if @before_save_callback
+        execute_query
+        @after_save_callback.() if @after_save_callback
+        @set.clear
+      end
+
+      self
+    end
+
+    def execute_query
+      if query = compose_insert_query
+        result_set = @connection.exec_query(query)
+        @result_sets.push(result_set) if @return_primary_keys
+      end
+    end
+
+    def compose_insert_query
+      sql = insert_sql_statement
+      @now = Time.now
+      rows = []
+
+      @set.each do |row|
+        values = []
+        @columns.zip(row) do |column, value|
+          value = @now if value == :__timestamp_placeholder
+
+          if ActiveRecord::VERSION::STRING >= "5.0.0"
+            value = @connection.lookup_cast_type_from_column(column, value) if column
+            values << @connection.quote(value)
+          else
+            values << @connection.quote(value, column)
+          end
+        end
+        rows << "(#{values.join(',')})"
+      end
+
+      if !rows.empty?
+        sql << rows.join(",")
+        sql << @statement_adapter.on_conflict_statement(@columns, ignore, update_duplicates)
+        sql << @statement_adapter.primary_key_return_statement(@primary_key) if @return_primary_keys
+        sql
+      else
+        false
+      end
+    end
+
+    def insert_sql_statement
+      insert_ignore = @ignore ? @statement_adapter.insert_ignore_statement : ''
+      "INSERT #{insert_ignore} INTO #{@table_name} (#{@column_names}) VALUES "
+    end
+  end
+end

data/test/bulk_insert/worker_test.rb

@@ -0,0 +1,459 @@
+require 'minitest/mock'
+require 'test_helper'
+
+class BulkInsertWorkerTest < ActiveSupport::TestCase
+  setup do
+    @insert = BulkInsert::Worker.new(
+      Testing.connection,
+      Testing.table_name,
+      'id',
+      %w(greeting age happy created_at updated_at color))
+    @now = Time.now
+  end
+
+  test "empty insert is not pending" do
+    assert_equal false, @insert.pending?
+  end
+
+  test "pending_count should describe size of pending set" do
+    assert_equal 0, @insert.pending_count
+    @insert.add ["Hello", 15, true, @now, @now]
+    assert_equal 1, @insert.pending_count
+  end
+
+  test "default set size" do
+    assert_equal 500, @insert.set_size
+  end
+
+  test "adding row to insert makes insert pending" do
+    @insert.add ["Hello", 15, true, @now, @now]
+    assert_equal true, @insert.pending?
+  end
+
+  test "add should default timestamp columns to current time" do
+    now = Time.now
+
+    @insert.add ["Hello", 15, true]
+    @insert.save!
+
+    record = Testing.first
+    assert_operator record.created_at, :>=, now
+    assert_operator record.updated_at, :>=, now
+  end
+
+  test "default timestamp columns should be equivalent for the entire batch" do
+    @insert.add ["Hello", 15, true]
+    @insert.add ["Howdy", 20, false]
+    @insert.save!
+
+    first, second = Testing.all
+    assert_equal first.created_at.to_f, second.created_at.to_f
+    assert_equal first.created_at.to_f, first.updated_at.to_f
+  end
+
+  test "add should use database default values when present" do
+    @insert.add greeting: "Hello", age: 20, happy: false
+    @insert.save!
+
+    record = Testing.first
+    assert_equal record.color, "chartreuse"
+  end
+
+  test "explicit nil should override defaults" do
+    @insert.add greeting: "Hello", age: 20, happy: false, color: nil
+    @insert.save!
+
+    record = Testing.first
+    assert_nil record.color
+  end
+
+  test "add should allow values given as Hash" do
+    @insert.add greeting: "Yo", age: 20, happy: false, created_at: @now, updated_at: @now
+    @insert.save!
+
+    record = Testing.first
+    assert_not_nil record
+    assert_equal "Yo", record.greeting
+    assert_equal 20, record.age
+    assert_equal false, record.happy?
+  end
+
+  test "add should save automatically when overflowing set size" do
+    @insert.set_size = 1
+    @insert.add ["Hello", 15, true, @now, @now]
+    @insert.add ["Yo", 20, false, @now, @now]
+    assert_equal 1, Testing.count
+    assert_equal "Hello", Testing.first.greeting
+  end
+
+  test "add_all should append all items to the set" do
+    @insert.add_all [
+      [ "Hello", 15, true ],
+      { greeting: "Hi", age: 55, happy: true }
+    ]
+    assert_equal 2, @insert.pending_count
+  end
+
+  test "save! makes insert not pending" do
+    @insert.add ["Hello", 15, true, @now, @now]
+    @insert.save!
+    assert_equal false, @insert.pending?
+  end
+
+  test "save! when not pending should do nothing" do
+    assert_no_difference 'Testing.count' do
+      @insert.save!
+    end
+  end
+
+  test "save! inserts pending records" do
+    @insert.add ["Yo", 15, false, @now, @now]
+    @insert.add ["Hello", 25, true, @now, @now]
+    @insert.save!
+
+    yo = Testing.find_by(greeting: 'Yo')
+    hello = Testing.find_by(greeting: 'Hello')
+
+    assert_not_nil yo
+    assert_equal 15, yo.age
+    assert_equal false, yo.happy?
+
+    assert_not_nil hello
+    assert_equal 25, hello.age
+    assert_equal true, hello.happy?
+  end
+
+  test "save! does not add to result sets when not returning primary keys" do
+    @insert.add greeting: "first"
+    @insert.add greeting: "second"
+    @insert.save!
+
+    assert_equal 0, @insert.result_sets.count
+  end
+
+
+  test "save! adds to result sets when returning primary keys" do
+    worker = BulkInsert::Worker.new(
+      Testing.connection,
+      Testing.table_name,
+      'id',
+      %w(greeting age happy created_at updated_at color),
+      500,
+      false,
+      false,
+      true
+    )
+
+    assert_no_difference -> { worker.result_sets.count } do
+      worker.save!
+    end
+
+    worker.add greeting: "first"
+    worker.add greeting: "second"
+    worker.save!
+    assert_equal 1, worker.result_sets.count
+
+    worker.add greeting: "third"
+    worker.add greeting: "fourth"
+    worker.save!
+    assert_equal 2, worker.result_sets.count
+  end
+
+  test "initialized with empty result sets array" do
+    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)
+    assert_empty new_worker.result_sets
+  end
+
+  test "save! calls the after_save handler" do
+    x = 41
+
+    @insert.after_save do
+      x += 1
+    end
+
+    @insert.add ["Yo", 15, false, @now, @now]
+    @insert.add ["Hello", 25, true, @now, @now]
+    @insert.save!
+
+    assert_equal 42, x
+  end
+
+  test "after_save stores a block as a proc" do
+    @insert.after_save do
+      "hello"
+    end
+
+    assert_equal "hello", @insert.after_save_callback.()
+  end
+
+  test "after_save_callback can be set as a proc" do
+    @insert.after_save_callback = -> do
+      "hello"
+    end
+
+    assert_equal "hello", @insert.after_save_callback.()
+  end
+
+  test "save! calls the before_save handler" do
+    x = 41
+
+    @insert.before_save do
+      x += 1
+    end
+
+    @insert.add ["Yo", 15, false, @now, @now]
+    @insert.add ["Hello", 25, true, @now, @now]
+    @insert.save!
+
+    assert_equal 42, x
+  end
+
+  test "before_save stores a block as a proc" do
+    @insert.before_save do
+      "hello"
+    end
+
+    assert_equal "hello", @insert.before_save_callback.()
+  end
+
+  test "before_save_callback can be set as a proc" do
+    @insert.before_save_callback = -> do
+      "hello"
+    end
+
+    assert_equal "hello", @insert.before_save_callback.()
+  end
+
+  test "before_save can manipulate the set" do
+    @insert.before_save do |set|
+      set.reject!{|row| row[0] == "Yo"}
+    end
+
+    @insert.add ["Yo", 15, false, @now, @now]
+    @insert.add ["Hello", 25, true, @now, @now]
+    @insert.save!
+
+    yo = Testing.find_by(greeting: 'Yo')
+    hello = Testing.find_by(greeting: 'Hello')
+
+    assert_nil yo
+    assert_not_nil hello
+  end
+
+  test "save! doesn't blow up if before_save emptying the set" do
+    @insert.before_save do |set|
+      set.clear
+    end
+
+    @insert.add ["Yo", 15, false, @now, @now]
+    @insert.add ["Hello", 25, true, @now, @now]
+    @insert.save!
+
+    yo = Testing.find_by(greeting: 'Yo')
+    hello = Testing.find_by(greeting: 'Hello')
+
+    assert_nil yo
+    assert_nil hello
+  end
+
+  test "adapter dependent default methods" do
+    assert_equal @insert.adapter_name, 'SQLite'
+    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,0,NULL,NULL,'chartreuse')"
+  end
+
+  test "adapter dependent mysql methods" do
+    connection = Testing.connection
+    connection.stub :adapter_name, 'MySQL' do
+      mysql_worker = BulkInsert::Worker.new(
+        connection,
+        Testing.table_name,
+        'id',
+        %w(greeting age happy created_at updated_at color),
+        500, # batch size
+        true  # ignore
+      )
+
+      assert_equal mysql_worker.adapter_name, 'MySQL'
+      assert_equal (mysql_worker.adapter_name == 'MySQL'), true
+      assert_equal mysql_worker.ignore, true
+      assert_equal ((mysql_worker.adapter_name == 'MySQL') & mysql_worker.ignore), true
+
+      mysql_worker.add ["Yo", 15, false, nil, nil]
+
+      assert_equal mysql_worker.compose_insert_query, "INSERT IGNORE INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,0,NULL,NULL,'chartreuse')"
+    end
+  end
+
+  test "adapter dependent mysql methods work for mysql2" do
+    connection = Testing.connection
+    connection.stub :adapter_name, 'Mysql2' do
+      mysql_worker = BulkInsert::Worker.new(
+        connection,
+        Testing.table_name,
+        'id',
+        %w(greeting age happy created_at updated_at color),
+        500, # batch size
+        true, # ignore
+        true) # update_duplicates
+
+      assert_equal mysql_worker.adapter_name, 'Mysql2'
+      assert mysql_worker.ignore
+
+      mysql_worker.add ["Yo", 15, false, nil, nil]
+
+      assert_equal mysql_worker.compose_insert_query, "INSERT IGNORE INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,0,NULL,NULL,'chartreuse') ON DUPLICATE KEY UPDATE `greeting`=VALUES(`greeting`), `age`=VALUES(`age`), `happy`=VALUES(`happy`), `created_at`=VALUES(`created_at`), `updated_at`=VALUES(`updated_at`), `color`=VALUES(`color`)"
+    end
+  end
+
+  test "adapter dependent Mysql2Spatial methods" do
+    connection = Testing.connection
+    connection.stub :adapter_name, 'Mysql2Spatial' do
+      mysql_worker = BulkInsert::Worker.new(
+        connection,
+        Testing.table_name,
+        'id',
+        %w(greeting age happy created_at updated_at color),
+        500, # batch size
+        true) # ignore
+
+      assert_equal mysql_worker.adapter_name, 'Mysql2Spatial'
+
+      mysql_worker.add ["Yo", 15, false, nil, nil]
+
+      assert_equal mysql_worker.compose_insert_query, "INSERT IGNORE INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,0,NULL,NULL,'chartreuse')"
+      end
+  end
+
+  test "adapter dependent postgresql methods" do
+    connection = Testing.connection
+    connection.stub :adapter_name, 'PostgreSQL' do
+      pgsql_worker = BulkInsert::Worker.new(
+        connection,
+        Testing.table_name,
+        'id',
+        %w(greeting age happy created_at updated_at color),
+        500, # batch size
+        true, # ignore
+        false, # update duplicates
+        true # return primary keys
+      )
+
+      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,0,NULL,NULL,'chartreuse') ON CONFLICT DO NOTHING RETURNING id"
+    end
+  end
+
+  test "adapter dependent postgresql methods (no ignore, no update_duplicates)" do
+    connection = Testing.connection
+    connection.stub :adapter_name, 'PostgreSQL' do
+      pgsql_worker = BulkInsert::Worker.new(
+        connection,
+        Testing.table_name,
+        'id',
+        %w(greeting age happy created_at updated_at color),
+        500, # batch size
+        false, # ignore
+        false, # update duplicates
+        true # return primary keys
+      )
+
+      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,0,NULL,NULL,'chartreuse') RETURNING id"
+    end
+  end
+
+  test "adapter dependent postgresql methods (with update_duplicates)" do
+    connection = Testing.connection
+    connection.stub :adapter_name, 'PostgreSQL' do
+      pgsql_worker = BulkInsert::Worker.new(
+        connection,
+        Testing.table_name,
+        'id',
+        %w(greeting age happy created_at updated_at color),
+        500, # batch size
+        false, # ignore
+        %w(greeting age happy), # update duplicates
+        true # return primary keys
+      )
+      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,0,NULL,NULL,'chartreuse') ON CONFLICT(greeting, age, happy) DO UPDATE SET greeting=EXCLUDED.greeting, age=EXCLUDED.age, happy=EXCLUDED.happy, created_at=EXCLUDED.created_at, updated_at=EXCLUDED.updated_at, color=EXCLUDED.color RETURNING id"
+    end
+  end
+
+  test "adapter dependent PostGIS methods" do
+    connection = Testing.connection
+    connection.stub :adapter_name, 'PostGIS' do
+      pgsql_worker = BulkInsert::Worker.new(
+        connection,
+        Testing.table_name,
+        'id',
+        %w(greeting age happy created_at updated_at color),
+        500, # batch size
+        true, # ignore
+        false, # update duplicates
+        true # return primary keys
+      )
+      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,0,NULL,NULL,'chartreuse') ON CONFLICT DO NOTHING RETURNING id"
+    end
+  end
+
+  test "adapter dependent sqlite3 methods (with lowercase adapter name)" do
+    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
+    sqlite_worker.adapter_name = 'sqlite3'
+    sqlite_worker.add ["Yo", 15, false, nil, nil]
+
+    assert_equal sqlite_worker.compose_insert_query, "INSERT OR IGNORE INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,0,NULL,NULL,'chartreuse')"
+  end
+
+  test "adapter dependent sqlite3 methods (with stylecase adapter name)" do
+    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
+    sqlite_worker.adapter_name = 'SQLite'
+    sqlite_worker.add ["Yo", 15, false, nil, nil]
+
+    assert_equal sqlite_worker.compose_insert_query, "INSERT OR IGNORE INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,0,NULL,NULL,'chartreuse')"
+  end
+
+  test "mysql adapter can update duplicates" do
+    connection = Testing.connection
+    connection.stub :adapter_name, 'MySQL' do
+      mysql_worker = BulkInsert::Worker.new(
+        connection,
+        Testing.table_name,
+        'id',
+        %w(greeting age happy created_at updated_at color),
+        500, # batch size
+        false, # ignore
+        true # update_duplicates
+      )
+      mysql_worker.add ["Yo", 15, false, nil, nil]
+
+      assert_equal mysql_worker.compose_insert_query, "INSERT  INTO \"testings\" (\"greeting\",\"age\",\"happy\",\"created_at\",\"updated_at\",\"color\") VALUES ('Yo',15,0,NULL,NULL,'chartreuse') ON DUPLICATE KEY UPDATE `greeting`=VALUES(`greeting`), `age`=VALUES(`age`), `happy`=VALUES(`happy`), `created_at`=VALUES(`created_at`), `updated_at`=VALUES(`updated_at`), `color`=VALUES(`color`)"
+    end
+  end
+end