Merge pull request #26 from willnet/upsert_all

Enable bulk upsert in SeedDo.seed

Author: willnet

.rubocop.yml

@@ -7,4 +7,6 @@ AllCops:
   NewCops: enable
   TargetRubyVersion: 3.2
 Performance:
-  Enabled: true
+  Enabled: true
+Metrics/ClassLength:
+  Enabled: false

README.md

@@ -142,6 +142,30 @@ Seed files can be run automatically using `rake db:seed_do`. There are two optio
 
 You can also do a similar thing in your code by calling `SeedDo.seed(fixture_paths, filter)`.
 
+## Bulk upsert
+
+If you load a large amount of seed data, you can enable bulk upsert mode:
+
+```ruby
+SeedDo.seed(bulk: true)
+```
+
+In bulk mode, SeedDo buffers seed data for each file and writes it with `upsert_all`, which can significantly reduce the number of queries.
+
+You can also control the batch size per `upsert_all` call. The default is `1000`.
+
+```ruby
+SeedDo.seed(bulk: { batch_size: 100 })
+```
+
+### Requirements and caveats
+
+- Bulk mode relies on Rails `upsert_all`
+- The seed constraints are passed to `unique_by`, so the related columns must have a unique index
+- Bulk mode is supported only on databases that support conflict targets, such as PostgreSQL and SQLite
+- Bulk mode is not available on MySQL
+- `SeedDo.seed` and `SeedDo.seed(bulk: true)` may not always produce the same record order, so verify the behavior carefully before enabling it in an existing production environment
+
 ## Disable output
 
 To disable output from Seed Do, set `SeedDo.quiet = true`.

lib/seed-do.rb

@@ -7,6 +7,7 @@ module SeedDo
   autoload :Seeder,                'seed-do/seeder'
   autoload :ActiveRecordExtension, 'seed-do/active_record_extension'
   autoload :BlockHash,             'seed-do/block_hash'
+  autoload :BulkSeeder,            'seed-do/bulk_seeder'
   autoload :Runner,                'seed-do/runner'
   autoload :Writer,                'seed-do/writer'
 
@@ -17,12 +18,20 @@ module SeedDo
   # plugin, SeedDo will set it to contain `Rails.root/db/fixtures` and
   # `Rails.root/db/fixtures/Rails.env`
   mattr_accessor :fixture_paths, default: ['db/fixtures']
-
   # Load seed data from files
   # @param [Array] fixture_paths The paths to look for seed files in
   # @param [Regexp] filter If given, only filenames matching this expression will be loaded
-  def self.seed(fixture_paths = SeedDo.fixture_paths, filter = nil)
-    Runner.new(fixture_paths, filter).run
+  # @param [Boolean] bulk If true, bulk insert/upsert will be used
+  def self.seed(fixture_paths = SeedDo.fixture_paths, filter = nil, bulk: false)
+    Runner.new(fixture_paths, filter, bulk: bulk).run
+  end
+
+  def self.current_seeder
+    @current_seeder || Seeder.new
+  end
+
+  def self.current_seeder=(seeder)
+    @current_seeder = seeder
   end
 end
 

lib/seed-do/active_record_extension.rb

@@ -29,7 +29,7 @@ module ActiveRecordExtension
     #       { :x => 5, :y => 9,  :name => "Office" }
     #     )
     def seed(*args, &block)
-      SeedDo::Seeder.new(self, *parse_seed_do_args(args, block)).seed
+      SeedDo.current_seeder.seed(self, *parse_seed_do_args(args, block))
     end
 
     # Has the same syntax as {#seed}, but if a record already exists with the same values for
@@ -40,8 +40,7 @@ def seed(*args, &block)
     #   Person.seed(:id, :id => 1, :name => "Bob") # => Name changed
     #   Person.seed_once(:id, :id => 1, :name => "Harry") # => Name *not* changed
     def seed_once(*args, &block)
-      constraints, data = parse_seed_do_args(args, block)
-      SeedDo::Seeder.new(self, constraints, data, insert_only: true).seed
+      SeedDo.current_seeder.seed_once(self, *parse_seed_do_args(args, block))
     end
 
     private

lib/seed-do/bulk_seeder.rb

@@ -0,0 +1,66 @@
+module SeedDo
+  # Buffers seeds during a file run and flushes them in bulk.
+  class BulkSeeder
+    attr_reader :buffer
+
+    def initialize(batch_size: 1000)
+      @batch_size = batch_size
+      @buffer = nil
+
+      validate_support!
+    end
+
+    def seed(model, constraints, data)
+      puts " - #{model} #{data.inspect}" unless SeedDo.quiet
+      buffer[:seed] << { model: model, constraints: constraints, data: data }
+    end
+
+    def seed_once(model, constraints, data)
+      puts " - #{model} #{data.inspect}" unless SeedDo.quiet
+      buffer[:seed_once] << { model: model, constraints: constraints, data: data }
+    end
+
+    def with_seed_file
+      @buffer = { seed: [], seed_once: [] }
+      yield
+      process_buffer
+    ensure
+      @buffer = nil
+    end
+
+    private
+
+    def validate_support!
+      return if ActiveRecord::Base.connection.supports_insert_conflict_target?
+
+      raise ArgumentError,
+            "Bulk mode is not supported for #{ActiveRecord::Base.connection.adapter_name}. " \
+            'The database does not support upsert operations with conflict targets. ' \
+            'Please use SeedDo.seed without the bulk option.'
+    end
+
+    def process_buffer
+      return unless buffer
+
+      buffer.each do |type, operations|
+        next if operations.empty?
+
+        operations.chunk { |operation| [operation[:model], operation[:constraints]] }
+                  .each do |(model, constraints), chunk|
+          flush_chunk(model, constraints, type, chunk)
+        end
+      end
+    end
+
+    def flush_chunk(model, constraints, type, chunk)
+      all_data = chunk.flat_map { |operation| operation[:data] }
+
+      options = { unique_by: constraints }
+      options[:on_duplicate] = :skip if type == :seed_once
+
+      all_data.each_slice(@batch_size) do |batch|
+        model.upsert_all(batch, **options)
+      end
+    end
+  end
+end

lib/seed-do/runner.rb

@@ -3,28 +3,39 @@
 
 module SeedDo
   # Runs seed files.
-  #
-  # It is not recommended to use this class directly. Instead, use {SeedDo.seed SeedDo.seed}, which creates
-  # an instead of {Runner} and calls {#run #run}.
+  # It is not recommended to use this class directly. Instead, use {SeedDo.seed SeedDo.seed}, which creates an instead of {Runner} and calls {#run #run}.
   #
   # @see SeedDo.seed SeedDo.seed
   class Runner
     # @param [Array<String>] fixture_paths The paths where fixtures are located. Will use
     #   `SeedDo.fixture_paths` if {nil}. If the argument is not an array, it will be wrapped by one.
     # @param [Regexp] filter If given, only seed files with a file name matching this pattern will
     #   be used
-    def initialize(fixture_paths = nil, filter = nil)
+    # @param [Boolean, Hash] bulk If true, use upsert_all to insert/update records in bulk.
+    #   If a hash, can include :batch_size (default: 1000) to control the number of records
+    #   per upsert_all call.
+    def initialize(fixture_paths = nil, filter = nil, bulk: false)
       @fixture_paths = Array.wrap(fixture_paths || SeedDo.fixture_paths)
       @filter        = filter
+      @seeder        = build_seeder(bulk)
     end
 
     # Run the seed files.
     def run
+      SeedDo.current_seeder = @seeder
       puts "\n== Filtering seed files against regexp: #{@filter.inspect}" if @filter && !SeedDo.quiet
 
-      filenames.each do |filename|
-        run_file(filename)
-      end
+      filenames.each { |filename| run_file(filename) }
+    ensure
+      SeedDo.current_seeder = nil
+    end
+
+    def seed(model, constraints, data)
+      @seeder.seed(model, constraints, data)
+    end
+
+    def seed_once(model, constraints, data)
+      @seeder.seed_once(model, constraints, data)
     end
 
     private
@@ -33,18 +44,28 @@ def run_file(filename)
       puts "\n== Seed from #{filename}" unless SeedDo.quiet
 
       ActiveRecord::Base.transaction do
-        open(filename) do |file|
-          chunked_ruby = +''
-          file.each_line do |line|
-            if line == "# BREAK EVAL\n"
-              eval(chunked_ruby)
-              chunked_ruby = +''
-            else
-              chunked_ruby << line
-            end
+        @seeder.with_seed_file { _run_file(filename) }
+      end
+    end
+
+    def build_seeder(bulk)
+      return SeedDo::BulkSeeder.new(batch_size: bulk.fetch(:batch_size, 1000)) if bulk.is_a?(Hash)
+
+      bulk ? SeedDo::BulkSeeder.new : SeedDo::Seeder.new
+    end
+
+    def _run_file(filename)
+      open(filename) do |file|
+        chunked_ruby = +''
+        file.each_line do |line|
+          if line == "# BREAK EVAL\n"
+            eval(chunked_ruby)
+            chunked_ruby = +''
+          else
+            chunked_ruby << line
           end
-          eval(chunked_ruby) unless chunked_ruby == ''
         end
+        eval(chunked_ruby) unless chunked_ruby == ''
       end
     end
 

lib/seed-do/seeder.rb

@@ -8,39 +8,38 @@ module SeedDo
   #
   # @see ActiveRecordExtension
   class Seeder
-    # @param [ActiveRecord::Base] model_class The model to be seeded
-    # @param [Array<Symbol>] constraints A list of attributes which identify a particular seed. If
-    #   a record with these attributes already exists then it will be updated rather than created.
-    # @param [Array<Hash>] data Each item in this array is a hash containing attributes for a
-    #   particular record.
-    # @param [Hash] options
-    # @option options [Boolean] :quiet (SeedDo.quiet) If true, output will be silenced
-    # @option options [Boolean] :insert_only (false) If true then existing records which match the
-    #   constraints will not be updated, even if the seed data has changed
-    def initialize(model_class, constraints, data, options = {})
+    def seed(model_class, constraints, data)
+      seed_records(model_class, constraints, data)
+    end
+
+    def seed_once(model_class, constraints, data)
+      seed_records(model_class, constraints, data, insert_only: true)
+    end
+
+    def with_seed_file
+      yield
+    end
+
+    private
+
+    # Insert/update the records as appropriate. Validation is skipped while saving.
+    # @return [Array<ActiveRecord::Base>] The records which have been seeded
+    def seed_records(model_class, constraints, data, options = {})
       @model_class = model_class
       @constraints = constraints.to_a.empty? ? [:id] : constraints
       @data        = data.to_a || []
       @options     = options.symbolize_keys
 
-      @options[:quiet] ||= SeedDo.quiet
-
       validate_constraints!
       validate_data!
-    end
 
-    # Insert/update the records as appropriate. Validation is skipped while saving.
-    # @return [Array<ActiveRecord::Base>] The records which have been seeded
-    def seed
       records = @model_class.transaction do
         @data.map { |record_data| seed_record(record_data.symbolize_keys) }
       end
       update_id_sequence
       records
     end
 
-    private
-
     def validate_constraints!
       unknown_columns = @constraints.map(&:to_s) - @model_class.column_names
       return if unknown_columns.empty?
@@ -62,7 +61,7 @@ def seed_record(data)
       record = find_or_initialize_record(data)
       return if @options[:insert_only] && !record.new_record?
 
-      puts " - #{@model_class} #{data.inspect}" unless @options[:quiet]
+      puts " - #{@model_class} #{data.inspect}" unless SeedDo.quiet
 
       record.assign_attributes(data)
       record.save(validate: false) || raise(ActiveRecord::RecordNotSaved, 'Record not saved!')

spec/bulk_seeder_spec.rb

@@ -0,0 +1,75 @@
+require 'spec_helper'
+
+describe SeedDo::BulkSeeder do
+  before(:all) do
+    skip 'Bulk mode is not supported on databases without insert conflict target support' unless ActiveRecord::Base.connection.supports_insert_conflict_target?
+  end
+
+  around do |example|
+    original_quiet = SeedDo.quiet
+    example.run
+    SeedDo.quiet = original_quiet
+  end
+
+  it 'clears the buffer after processing' do
+    bulk_seeder = described_class.new(batch_size: 2)
+
+    bulk_seeder.with_seed_file do
+      bulk_seeder.seed(BulkSeededModel, [:title], [{ title: 'Bulk 1', login: 'bulk1' }])
+      expect(bulk_seeder.buffer).not_to be_nil
+    end
+
+    expect(bulk_seeder.buffer).to be_nil
+  end
+
+  it 'separates flushes by operation type' do
+    bulk_seeder = described_class.new(batch_size: 10)
+
+    expect(BulkSeededModel).to receive(:upsert_all).twice.and_call_original
+
+    bulk_seeder.with_seed_file do
+      bulk_seeder.seed(BulkSeededModel, [:title], [{ title: 'Bulk 1', login: 'bulk1' }])
+      bulk_seeder.seed(BulkSeededModel, [:title], [{ title: 'Bulk 2', login: 'bulk2' }])
+      bulk_seeder.seed_once(BulkSeededModel, [:title], [{ title: 'Bulk 3', login: 'bulk3' }])
+    end
+  end
+
+  it 'logs buffered seed operations when quiet is false' do
+    SeedDo.quiet = false
+    bulk_seeder = described_class.new(batch_size: 10)
+
+    output = capture_stdout do
+      bulk_seeder.with_seed_file do
+        bulk_seeder.seed(BulkSeededModel, [:title], [{ title: 'Bulk 1', login: 'bulk1' }])
+        bulk_seeder.seed_once(BulkSeededModel, [:title], [{ title: 'Bulk 2', login: 'bulk2' }])
+      end
+    end
+
+    expect(output).to include(' - BulkSeededModel [')
+    expect(output).to include('Bulk 1')
+    expect(output).to include('bulk1')
+    expect(output).to include('Bulk 2')
+    expect(output).to include('bulk2')
+    expect(output.lines.count).to eq(2)
+  end
+
+  it 'does not log buffered seed operations when quiet is true' do
+    SeedDo.quiet = true
+    bulk_seeder = described_class.new(batch_size: 10)
+
+    expect do
+      bulk_seeder.with_seed_file do
+        bulk_seeder.seed(BulkSeededModel, [:title], [{ title: 'Bulk 1', login: 'bulk1' }])
+      end
+    end.not_to output.to_stdout
+  end
+
+  def capture_stdout
+    original_stdout = $stdout
+    $stdout = StringIO.new
+    yield
+    $stdout.string
+  ensure
+    $stdout = original_stdout
+  end
+end