active_record_data_loader 1.2.0 → 1.3.0

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,19 +33,23 @@ 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)
           @settings.queries[klass].call.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,16 +2,18 @@
 
 module ActiveRecordDataLoader
   class BulkInsertStrategy
-    def initialize(data_generator, output_adapter)
+    def initialize(data_generator, file_adapter)
       @data_generator = data_generator
-      @output_adapter = output_adapter
+      @file_adapter = file_adapter
     end
 
     def load_batch(row_numbers, connection)
-      output_adapter.insert(connection: connection, command: <<~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
@@ -24,7 +26,11 @@ module ActiveRecordDataLoader
 
     private
 
-    attr_reader :data_generator, :output_adapter
+    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)
@@ -39,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

@@ -3,7 +3,7 @@
 module ActiveRecordDataLoader
   class Configuration
     attr_accessor :connection_factory, :default_batch_size, :default_row_count,
-                  :logger, :statement_timeout
+                  :logger, :max_duplicate_retries, :raise_on_duplicates, :statement_timeout
     attr_reader :output
 
     def initialize(
@@ -12,51 +12,34 @@ module ActiveRecordDataLoader
       logger: nil,
       statement_timeout: "2min",
       connection_factory: -> { ::ActiveRecord::Base.connection },
-      output: :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 || { type: :connection })
-    end
-
-    def output_adapter
-      if output.fetch(:type) == :file
-        ActiveRecordDataLoader::FileOutputAdapter.new(output)
-      else
-        ActiveRecordDataLoader::ConnectionOutputAdapter.new
-      end
-    end
-
-    def connection_handler
-      ActiveRecordDataLoader::ConnectionHandler.new(
-        connection_factory: connection_factory,
-        statement_timeout: statement_timeout,
-        output_adapter: output_adapter
-      )
+      @output = validate_output(output)
     end
 
     private
 
-    OUTPUT_OPTIONS_BY_TYPE = { connection: %i[type], file: %i[type filename] }.freeze
-
     def validate_output(output)
-      if %i[file connection].include?(output)
-        { type: output }
-      elsif output.is_a?(Hash)
-        raise "The output hash must contain a :type key with either :connection or :file" \
-          unless %i[file connection].include?(output[:type])
-
-        output.slice(*OUTPUT_OPTIONS_BY_TYPE[output[:type]])
+      if output.to_s.blank?
+        nil
+      elsif output.is_a?(String)
+        output
       else
-        raise "The output configuration parameter must be either a symbol for :connection or :file, "\
-              "or a hash with more detailed output options."
+        raise "The output configuration parameter must be a filename meant to be the "\
+              "target for the SQL script"
       end
     end
 

data/lib/active_record_data_loader/connection_handler.rb

@@ -2,19 +2,18 @@
 
 module ActiveRecordDataLoader
   class ConnectionHandler
-    def initialize(connection_factory:, statement_timeout:, output_adapter:)
+    def initialize(connection_factory:, statement_timeout:)
       @connection_factory = connection_factory
       @statement_timeout = statement_timeout
-      @output_adapter = output_adapter
+      cache_facts
     end
 
     def with_connection
-      connection = open_connection
-      if postgres?(connection)
-        original_timeout = retrieve_statement_timeout(connection)
-        update_statement_timeout(connection, statement_timeout)
+      connection = connection_factory.call
+      if supports_timeout?
+        connection.execute(timeout_set_command)
         yield connection
-        update_statement_timeout(connection, original_timeout)
+        connection.execute(reset_timeout_command)
       else
         yield connection
       end
@@ -22,53 +21,32 @@ module ActiveRecordDataLoader
       connection&.close
     end
 
-    # When the output is going to a script file, there are two places to update the
-    # statement_timeout. The connection itself needs to have the timeout updated
-    # because we are reading data from the connection to come up with related data
-    # while generating the data. Also, the final SQL script file needs the timeout
-    # updated so that when those \COPY commands are executed they have the higher
-    # timeout as well.
-    def with_statement_timeout_for_output
-      return yield unless output_adapter.needs_timeout_output?
-
-      original_timeout = begin
-        connection = open_connection
-        retrieve_statement_timeout(connection) if postgres?(connection)
-      ensure
-        connection&.close
-      end
-
-      if original_timeout
-        output_adapter.execute(statement_timeout_set_command(statement_timeout))
-        yield
-        output_adapter.execute(statement_timeout_set_command(original_timeout))
-      else
-        yield
-      end
+    def supports_timeout?
+      @supports_timeout
     end
 
-    private
-
-    attr_reader :connection_factory, :statement_timeout, :output_adapter
-
-    def retrieve_statement_timeout(connection)
-      connection.execute("SHOW statement_timeout").first["statement_timeout"]
+    def supports_copy?
+      @supports_copy
     end
 
-    def update_statement_timeout(connection, timeout)
-      connection.execute(statement_timeout_set_command(timeout))
+    def timeout_set_command
+      "SET statement_timeout = \"#{statement_timeout}\""
     end
 
-    def statement_timeout_set_command(timeout)
-      "SET statement_timeout = \"#{timeout}\""
+    def reset_timeout_command
+      "RESET statement_timeout"
     end
 
-    def open_connection
-      connection_factory.call
-    end
+    private
 
-    def postgres?(connection)
-      connection.adapter_name.downcase.to_sym == :postgresql
+    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,17 +2,24 @@
 
 module ActiveRecordDataLoader
   class CopyStrategy
-    def initialize(data_generator, output_adapter)
+    def initialize(data_generator, file_adapter)
       @data_generator = data_generator
-      @output_adapter = output_adapter
+      @file_adapter = file_adapter
     end
 
     def load_batch(row_numbers, connection)
-      output_adapter.copy(
+      data = csv_rows(row_numbers, connection)
+      copy(
         connection: connection,
         table: table_name_for_copy(connection),
         columns: columns_for_copy(connection),
-        data: csv_rows(row_numbers, 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
@@ -27,12 +34,19 @@ module ActiveRecordDataLoader
 
     private
 
-    attr_reader :data_generator, :output_adapter
+    attr_reader :data_generator, :file_adapter
+
+    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
+        data_generator.generate_row(i)&.map { |d| quote_data(d, connection) }&.join(",")
+      end.compact
     end
 
     def table_name_for_copy(connection)

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

@@ -2,38 +2,46 @@
 
 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 needs_timeout_output?
-      true
-    end
-
-    def copy(connection:, table:, columns:, data:, row_numbers:)
+    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.open(filename, "a") do |file|
         file.puts("\\COPY #{table} (#{columns}) FROM '#{data_filename}' WITH (FORMAT CSV);")
       end
     end
 
-    def insert(connection:, command:)
-      execute(command)
+    def insert(command)
+      write_command(command)
     end
 
-    def execute(command)
-      File.open(@filename, "a") { |f| f.puts("#{command.gsub("\n", ' ')};") }
+    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"
+        path,
+        "#{file_basename}_#{table.gsub(/"/, '')}_rows_#{row_numbers[0]}_to_#{row_numbers[-1]}.csv"
       )
     end
   end

data/lib/active_record_data_loader/loader.rb

@@ -1,76 +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,
-          connection_handler: configuration.connection_handler,
-          strategy: strategy_class(configuration.connection_factory).new(
-            data_generator,
-            configuration.output_adapter
-          )
-        ).load_data(batch_size, total_rows)
-      end
+    def initialize(configuration, definition)
+      @configuration = configuration
+      @definition = definition
+    end
 
-      private
+    def load_data
+      ActiveRecordDataLoader::ActiveRecord::PerRowValueCache.clear
 
-      def strategy_class(connection_factory)
-        if connection_factory.call.raw_connection.respond_to?(:copy_data)
-          ActiveRecordDataLoader::CopyStrategy
-        else
-          ActiveRecordDataLoader::BulkInsertStrategy
-        end
+      file_adapter_class.with_output_options(file_adapter_options) do |file_adapter|
+        definition.models.map { |m| load_model(m, file_adapter) }
       end
     end
 
-    def initialize(logger:, connection_handler:, strategy:)
-      @logger = logger
-      @connection_handler = connection_handler
-      @strategy = strategy
-    end
+    private
 
-    def load_data(batch_size, total_rows)
-      batch_count = (total_rows / batch_size.to_f).ceil
+    attr_reader :definition, :configuration
 
-      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)."
+    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
       )
-      total_time = Benchmark.realtime do
-        load_in_batches(batch_size, total_rows, batch_count)
-      end
-      logger.info(
-        "[ActiveRecordDataLoader] "\
-        "Completed loading #{total_rows} row(s) into '#{strategy.table_name}' "\
-        "in #{total_time} seconds."
+    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
 
-    private
+    def file_adapter_class
+      if configuration.output.present?
+        ActiveRecordDataLoader::FileOutputAdapter
+      else
+        ActiveRecordDataLoader::NullOutputAdapter
+      end
+    end
 
-    attr_reader :strategy, :connection_handler, :logger
+    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
 
-    def load_in_batches(batch_size, total_rows, batch_count)
-      connection_handler.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) }
+      timeout_commands.merge(filename: configuration.output)
+    end
 
-          logger.debug(
-            "[ActiveRecordDataLoader] "\
-            "Completed batch #{i + 1}/#{batch_count}, #{row_numbers.count} row(s) in #{time} seconds"
-          )
-        end
-      end
+    def strategy_class
+      @strategy_class ||= if connection_handler.supports_copy?
+                            ActiveRecordDataLoader::CopyStrategy
+                          else
+                            ActiveRecordDataLoader::BulkInsertStrategy
+                          end
+    end
+
+    def connection_handler
+      @connection_handler ||= ActiveRecordDataLoader::ConnectionHandler.new(
+        connection_factory: configuration.connection_factory,
+        statement_timeout: configuration.statement_timeout
+      )
     end
   end
 end