active_record_data_loader 1.0.2 → 1.3.1

data/lib/active_record_data_loader/active_record/{polymorphic_belongs_to_configuration.rb → polymorphic_belongs_to_data_provider.rb}

@@ -2,20 +2,21 @@
 
 module ActiveRecordDataLoader
   module ActiveRecord
-    class PolymorphicBelongsToConfiguration
-      def self.config_for(polymorphic_settings:)
+    class PolymorphicBelongsToDataProvider
+      def self.provider_for(polymorphic_settings:, strategy: :random)
         ar_association = polymorphic_settings.model_class.reflect_on_association(
           polymorphic_settings.name
         )
         raise "#{name} only supports polymorphic associations" unless ar_association.polymorphic?
 
-        new(polymorphic_settings, ar_association).polymorphic_config
+        new(polymorphic_settings, ar_association, strategy).polymorphic_config
       end
 
-      def initialize(settings, ar_association)
+      def initialize(settings, ar_association, strategy)
         @settings = settings
         @ar_association = ar_association
         @model_count = settings.weighted_models.size
+        @strategy = strategy
       end
 
       def polymorphic_config
@@ -32,21 +33,25 @@ module ActiveRecordDataLoader
       end
 
       def foreign_key(row_number)
-        possible_values[row_number % @model_count][1].sample
+        possible_values[row_number % @model_count][1].next
       end
 
       def possible_values
         @possible_values ||= begin
           values = @settings.models.keys.map do |klass|
-            [klass.name, base_query(klass).pluck(klass.primary_key).to_a]
+            [klass.name, values_query(klass)]
           end.to_h
 
           @settings.weighted_models.map { |klass| [klass.name, values[klass.name]] }
         end
       end
 
+      def values_query(klass)
+        List.for(base_query(klass).pluck(klass.primary_key), strategy: @strategy)
+      end
+
       def base_query(klass)
-        if @settings.queries[klass]&.respond_to?(:call)
+        if @settings.queries[klass].respond_to?(:call)
           @settings.queries[klass].call.all
         else
           klass.all

data/lib/active_record_data_loader/active_record/unique_index_tracker.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+module ActiveRecordDataLoader
+  module ActiveRecord
+    class UniqueIndexTracker
+      Index = Struct.new(:name, :columns, :column_indexes, keyword_init: true)
+
+      def initialize(model:, connection_factory:)
+        @model = model
+        @table = model.table_name
+        @unique_indexes = []
+        @unique_values_used = {}
+        find_unique_indexes(connection_factory)
+      end
+
+      def map_indexed_columns(column_list)
+        @unique_indexes = @raw_unique_indexes.map do |index|
+          @unique_values_used[index.name] = Set.new
+          columns = index.columns.map(&:to_sym)
+          Index.new(
+            name: index.name,
+            columns: columns,
+            column_indexes: columns.map { |c| column_list.find_index(c) }
+          )
+        end
+      end
+
+      def repeating_unique_values?(row)
+        @unique_indexes.map do |index|
+          values = index.column_indexes.map { |i| row[i] }
+          @unique_values_used.fetch(index.name).include?(values)
+        end.any?
+      end
+
+      def capture_unique_values(row)
+        return unless row.present?
+
+        @unique_indexes.each do |index|
+          values = index.column_indexes.map { |i| row[i] }
+          @unique_values_used.fetch(index.name) << values
+        end
+        row
+      end
+
+      def contained_in_index?(ar_column)
+        target_column = if @model.reflect_on_association(ar_column.name)&.belongs_to?
+                          ar_column.join_foreign_key.to_sym
+                        else
+                          ar_column.name.to_sym
+                        end
+
+        @raw_unique_indexes.flat_map { |i| i.columns.map(&:to_sym) }.include?(target_column)
+      end
+
+      private
+
+      attr_reader :table
+
+      def find_unique_indexes(connection_factory)
+        connection = connection_factory.call
+        @raw_unique_indexes = connection.indexes(table).select(&:unique)
+      ensure
+        connection&.close
+      end
+    end
+  end
+end

data/lib/active_record_data_loader/bulk_insert_strategy.rb

@@ -2,15 +2,18 @@
 
 module ActiveRecordDataLoader
   class BulkInsertStrategy
-    def initialize(data_generator)
+    def initialize(data_generator, file_adapter)
       @data_generator = data_generator
+      @file_adapter = file_adapter
     end
 
     def load_batch(row_numbers, connection)
-      connection.insert(<<~SQL)
+      command = <<~SQL
         INSERT INTO #{quoted_table_name(connection)} (#{column_list(connection)})
         VALUES #{values(row_numbers, connection)}
       SQL
+      insert(connection: connection, command: command)
+      file_adapter.insert(command)
     end
 
     def table_name
@@ -23,7 +26,11 @@ module ActiveRecordDataLoader
 
     private
 
-    attr_reader :data_generator
+    attr_reader :data_generator, :file_adapter
+
+    def insert(connection:, command:)
+      connection.insert(command)
+    end
 
     def quoted_table_name(connection)
       @quoted_table_name ||= connection.quote_table_name(data_generator.table)
@@ -38,15 +45,16 @@ module ActiveRecordDataLoader
 
     def values(row_numbers, connection)
       row_numbers
-        .map { |i| "(#{row_values(i, connection)})" }
+        .map { |i| row_values(i, connection) }
+        .compact
         .join(",")
     end
 
     def row_values(row_number, connection)
-      data_generator
-        .generate_row(row_number)
-        .map { |v| connection.quote(v) }
-        .join(",")
+      row = data_generator.generate_row(row_number)
+      return unless row.present?
+
+      "(#{row.map { |v| connection.quote(v) }.join(',')})"
     end
   end
 end

data/lib/active_record_data_loader/configuration.rb

@@ -2,29 +2,52 @@
 
 module ActiveRecordDataLoader
   class Configuration
-    attr_accessor :default_batch_size, :default_row_count, :logger, :statement_timeout, :connection_factory
+    attr_accessor :connection_factory, :default_batch_size, :default_row_count,
+                  :logger, :max_duplicate_retries, :raise_on_duplicates, :statement_timeout
+    attr_reader :output
 
     def initialize(
       default_batch_size: 100_000,
       default_row_count: 1,
       logger: nil,
       statement_timeout: "2min",
-      connection_factory: -> { ::ActiveRecord::Base.connection }
+      connection_factory: -> { ::ActiveRecord::Base.connection },
+      raise_on_duplicates: false,
+      max_duplicate_retries: 5,
+      output: nil
     )
       @default_batch_size = default_batch_size
       @default_row_count = default_row_count
       @logger = logger || default_logger
       @statement_timeout = statement_timeout
       @connection_factory = connection_factory
+      @raise_on_duplicates = raise_on_duplicates
+      @max_duplicate_retries = max_duplicate_retries
+      self.output = output
+    end
+
+    def output=(output)
+      @output = validate_output(output)
     end
 
     private
 
+    def validate_output(output)
+      if output.to_s.blank?
+        nil
+      elsif output.is_a?(String)
+        output
+      else
+        raise "The output configuration parameter must be a filename meant to be the "\
+              "target for the SQL script"
+      end
+    end
+
     def default_logger
       if defined?(Rails) && Rails.respond_to?(:logger)
         Rails.logger
       else
-        Logger.new(STDOUT, level: :info)
+        Logger.new($stdout, level: :info)
       end
     end
   end

data/lib/active_record_data_loader/connection_handler.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module ActiveRecordDataLoader
+  class ConnectionHandler
+    def initialize(connection_factory:, statement_timeout:)
+      @connection_factory = connection_factory
+      @statement_timeout = statement_timeout
+      cache_facts
+    end
+
+    def with_connection
+      connection = connection_factory.call
+      if supports_timeout?
+        connection.execute(timeout_set_command)
+        yield connection
+        connection.execute(reset_timeout_command)
+      else
+        yield connection
+      end
+    ensure
+      connection&.close
+    end
+
+    def supports_timeout?
+      @supports_timeout
+    end
+
+    def supports_copy?
+      @supports_copy
+    end
+
+    def timeout_set_command
+      "SET statement_timeout = \"#{statement_timeout}\""
+    end
+
+    def reset_timeout_command
+      "RESET statement_timeout"
+    end
+
+    private
+
+    attr_reader :connection_factory, :statement_timeout
+
+    def cache_facts
+      connection = connection_factory.call
+      @supports_timeout = connection.adapter_name.downcase.to_sym == :postgresql
+      @supports_copy = connection.raw_connection.respond_to?(:copy_data)
+    ensure
+      connection&.close
+    end
+  end
+end

data/lib/active_record_data_loader/copy_strategy.rb

@@ -2,15 +2,26 @@
 
 module ActiveRecordDataLoader
   class CopyStrategy
-    def initialize(data_generator)
+    def initialize(data_generator, file_adapter)
       @data_generator = data_generator
+      @file_adapter = file_adapter
     end
 
     def load_batch(row_numbers, connection)
-      csv_data = csv_data_batch(row_numbers, connection)
-
-      raw_connection = connection.raw_connection
-      raw_connection.copy_data(copy_command(connection)) { raw_connection.put_copy_data(csv_data) }
+      data = csv_rows(row_numbers, connection)
+      copy(
+        connection: connection,
+        table: table_name_for_copy(connection),
+        columns: columns_for_copy(connection),
+        data: data,
+        row_numbers: row_numbers
+      )
+      file_adapter.copy(
+        table: table_name_for_copy(connection),
+        columns: columns_for_copy(connection),
+        data: data,
+        row_numbers: row_numbers
+      )
     end
 
     def table_name
@@ -23,29 +34,32 @@ module ActiveRecordDataLoader
 
     private
 
-    attr_reader :data_generator
+    attr_reader :data_generator, :file_adapter
 
-    def csv_data_batch(row_numbers, connection)
-      row_numbers.map do |i|
-        data_generator.generate_row(i).map { |d| quote_data(d, connection) }.join(",")
-      end.join("\n")
-    end
-
-    def copy_command(connection)
-      @copy_command ||= begin
-        quoted_table_name = connection.quote_table_name(data_generator.table)
-        columns = data_generator
-                  .column_list
-                  .map { |c| connection.quote_column_name(c) }
-                  .join(", ")
-
-        <<~SQL
-          COPY #{quoted_table_name} (#{columns})
-            FROM STDIN WITH (FORMAT CSV)
-        SQL
+    def copy(connection:, table:, columns:, data:, row_numbers:)
+      raw_connection = connection.raw_connection
+      raw_connection.copy_data("COPY #{table} (#{columns}) FROM STDIN WITH (FORMAT CSV)") do
+        raw_connection.put_copy_data(data.join("\n"))
       end
     end
 
+    def csv_rows(row_numbers, connection)
+      row_numbers.map do |i|
+        data_generator.generate_row(i)&.map { |d| quote_data(d, connection) }&.join(",")
+      end.compact
+    end
+
+    def table_name_for_copy(connection)
+      @table_name_for_copy ||= connection.quote_table_name(data_generator.table)
+    end
+
+    def columns_for_copy(connection)
+      @columns_for_copy ||= data_generator
+                            .column_list
+                            .map { |c| connection.quote_column_name(c) }
+                            .join(", ")
+    end
+
     def quote_data(data, connection)
       return if data.nil?
 

data/lib/active_record_data_loader/data_faker.rb

@@ -13,16 +13,24 @@ module ActiveRecordDataLoader
 
       def adapter
         @adapter ||=
-          if Gem.loaded_specs.key?("ffaker")
-            require "ffaker"
+          if can_use?("ffaker", "2.1.0")
             FFakerGemAdapter.new
-          elsif Gem.loaded_specs.key?("faker")
-            require "faker"
+          elsif can_use?("faker", "1.9.3")
             FakerGemAdapter.new
           else
             NoGemAdapter.new
           end
       end
+
+      def can_use?(gem, min_version)
+        gemspec = Gem.loaded_specs[gem]
+        return false unless gemspec.present? && gemspec.version >= Gem::Version.new(min_version)
+
+        require gem
+        true
+      rescue LoadError
+        false
+      end
     end
 
     class FFakerGemAdapter

data/lib/active_record_data_loader/dsl/model.rb

@@ -3,13 +3,16 @@
 module ActiveRecordDataLoader
   module Dsl
     class Model
-      attr_reader :klass, :columns, :row_count, :polymorphic_associations, :belongs_to_associations
+      attr_reader :klass, :columns, :row_count, :polymorphic_associations, :belongs_to_associations,
+                  :raise_on_duplicates_flag
 
       def initialize(klass:, configuration:)
         @klass = klass
         @columns = {}
         @row_count = configuration.default_row_count
         @batch_size = configuration.default_batch_size
+        @raise_on_duplicates_flag = configuration.raise_on_duplicates
+        @max_duplicate_retries = configuration.max_duplicate_retries
         @polymorphic_associations = []
         @belongs_to_associations = []
       end
@@ -22,6 +25,20 @@ module ActiveRecordDataLoader
         @batch_size = (size || @batch_size)
       end
 
+      def raise_on_duplicates
+        @raise_on_duplicates_flag = true
+      end
+
+      def do_not_raise_on_duplicates
+        @raise_on_duplicates_flag = false
+      end
+
+      def max_duplicate_retries(retries = nil)
+        return @max_duplicate_retries if retries.nil?
+
+        @max_duplicate_retries = retries
+      end
+
       def column(name, func)
         @columns[name.to_sym] = func
       end
@@ -32,7 +49,7 @@ module ActiveRecordDataLoader
         ).tap { |a| block.call(a) }
       end
 
-      def belongs_to(assoc_name, eligible_set:)
+      def belongs_to(assoc_name, eligible_set: nil)
         @belongs_to_associations << BelongsToAssociation.new(@klass, assoc_name, eligible_set)
       end
     end

data/lib/active_record_data_loader/errors.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module ActiveRecordDataLoader
+  class DuplicateKeyError < StandardError; end
+end

data/lib/active_record_data_loader/file_output_adapter.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module ActiveRecordDataLoader
+  class FileOutputAdapter
+    def self.with_output_options(options)
+      adapter = new(options)
+      pre_command = options[:pre_command]
+      adapter.write_command(pre_command) if pre_command
+      yield adapter
+      post_command = options[:post_command]
+      adapter.write_command(post_command) if post_command
+    end
+
+    def initialize(options)
+      @filename = options.fetch(:filename, "active_record_data_loader_script.sql")
+      @file_basename = File.basename(@filename, File.extname(@filename))
+      @path = File.expand_path(File.dirname(@filename))
+      File.open(@filename, File::TRUNC) if File.exist?(@filename)
+    end
+
+    def copy(table:, columns:, data:, row_numbers:)
+      data_filename = data_filename(table, row_numbers)
+      File.open(data_filename, "w") { |f| f.puts(data) }
+      File.open(filename, "a") do |file|
+        file.puts("\\COPY #{table} (#{columns}) FROM '#{data_filename}' WITH (FORMAT CSV);")
+      end
+    end
+
+    def insert(command)
+      write_command(command)
+    end
+
+    def write_command(command)
+      File.open(filename, "a") { |f| f.puts("#{command.gsub("\n", ' ')};") }
+    end
+
+    private
+
+    attr_reader :filename, :path, :file_basename
+
+    def data_filename(table, row_numbers)
+      File.join(
+        path,
+        "#{file_basename}_#{table.gsub(/"/, '')}_rows_#{row_numbers[0]}_to_#{row_numbers[-1]}.csv"
+      )
+    end
+  end
+end

data/lib/active_record_data_loader/loader.rb

@@ -1,98 +1,82 @@
 # frozen_string_literal: true
 
-require "benchmark"
-
 module ActiveRecordDataLoader
   class Loader
-    class << self
-      def load_data(
-        data_generator:,
-        total_rows:,
-        batch_size:,
-        configuration:
-      )
-        new(
-          logger: configuration.logger,
-          statement_timeout: configuration.statement_timeout,
-          strategy: strategy_class(configuration.connection_factory).new(data_generator),
-          connection_factory: configuration.connection_factory
-        ).load_data(batch_size, total_rows)
-      end
-
-      private
-
-      def strategy_class(connection_factory)
-        if connection_factory.call.raw_connection.respond_to?(:copy_data)
-          ActiveRecordDataLoader::CopyStrategy
-        else
-          ActiveRecordDataLoader::BulkInsertStrategy
-        end
-      end
+    def initialize(configuration, definition)
+      @configuration = configuration
+      @definition = definition
     end
 
-    def initialize(logger:, statement_timeout:, strategy:, connection_factory:)
-      @logger = logger
-      @strategy = strategy
-      @statement_timeout = statement_timeout
-      @connection_factory = connection_factory
-    end
+    def load_data
+      ActiveRecordDataLoader::ActiveRecord::PerRowValueCache.clear
 
-    def load_data(batch_size, total_rows)
-      batch_count = (total_rows / batch_size.to_f).ceil
-
-      logger.info(
-        "[ActiveRecordDataLoader] "\
-        "Loading #{total_rows} row(s) into '#{strategy.table_name}' via #{strategy.name}. "\
-        "#{batch_size} row(s) per batch, #{batch_count} batch(es)."
-      )
-      total_time = Benchmark.realtime do
-        load_in_batches(batch_size, total_rows, batch_count)
+      file_adapter_class.with_output_options(file_adapter_options) do |file_adapter|
+        definition.models.map { |m| load_model(m, file_adapter) }
       end
-      logger.info(
-        "[ActiveRecordDataLoader] "\
-        "Completed loading #{total_rows} row(s) into '#{strategy.table_name}' "\
-        "in #{total_time} seconds."
-      )
     end
 
     private
 
-    attr_reader :strategy, :statement_timeout, :logger, :connection_factory
+    attr_reader :definition, :configuration
 
-    def load_in_batches(batch_size, total_rows, batch_count)
-      with_connection do |connection|
-        total_rows.times.each_slice(batch_size).with_index do |row_numbers, i|
-          time = Benchmark.realtime { strategy.load_batch(row_numbers, connection) }
+    def load_model(model, file_adapter)
+      ActiveRecordDataLoader::TableLoader.load_data(
+        batch_size: model.batch_size,
+        total_rows: model.row_count,
+        connection_handler: connection_handler,
+        strategy: strategy_class.new(generator(model), file_adapter),
+        logger: configuration.logger
+      )
+    end
 
-          logger.debug(
-            "[ActiveRecordDataLoader] "\
-            "Completed batch #{i + 1}/#{batch_count}, #{row_numbers.count} row(s) in #{time} seconds"
-          )
-        end
-      end
+    def generator(model)
+      ActiveRecordDataLoader::ActiveRecord::ModelDataGenerator.new(
+        model: model.klass,
+        column_settings: model.columns,
+        polymorphic_settings: model.polymorphic_associations,
+        belongs_to_settings: model.belongs_to_associations,
+        connection_factory: configuration.connection_factory,
+        raise_on_duplicates: model.raise_on_duplicates_flag,
+        max_duplicate_retries: model.max_duplicate_retries,
+        logger: configuration.logger
+      )
     end
 
-    def with_connection
-      if connection.adapter_name.downcase.to_sym == :postgresql
-        original_timeout = retrieve_statement_timeout
-        update_statement_timeout(statement_timeout)
-        yield connection
-        update_statement_timeout(original_timeout)
+    def file_adapter_class
+      if configuration.output.present?
+        ActiveRecordDataLoader::FileOutputAdapter
       else
-        yield connection
+        ActiveRecordDataLoader::NullOutputAdapter
       end
     end
 
-    def retrieve_statement_timeout
-      connection.execute("SHOW statement_timeout").first["statement_timeout"]
+    def file_adapter_options
+      timeout_commands =
+        if connection_handler.supports_timeout?
+          {
+            pre_command: connection_handler.timeout_set_command,
+            post_command: connection_handler.reset_timeout_command,
+          }
+        else
+          {}
+        end
+
+      timeout_commands.merge(filename: configuration.output)
     end
 
-    def update_statement_timeout(timeout)
-      connection.execute("SET statement_timeout = \"#{timeout}\"")
+    def strategy_class
+      @strategy_class ||= if connection_handler.supports_copy?
+                            ActiveRecordDataLoader::CopyStrategy
+                          else
+                            ActiveRecordDataLoader::BulkInsertStrategy
+                          end
     end
 
-    def connection
-      connection_factory.call
+    def connection_handler
+      @connection_handler ||= ActiveRecordDataLoader::ConnectionHandler.new(
+        connection_factory: configuration.connection_factory,
+        statement_timeout: configuration.statement_timeout
+      )
     end
   end
 end

data/lib/active_record_data_loader/null_output_adapter.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module ActiveRecordDataLoader
+  class NullOutputAdapter
+    def self.with_output_options(_options)
+      yield new
+    end
+
+    def copy(table:, columns:, data:, row_numbers:); end
+
+    def insert(command); end
+
+    def write_command(command); end
+  end
+end