active_record_data_loader 1.0.1 → 1.3.0

data/lib/active_record_data_loader/active_record/model_data_generator.rb

@@ -5,12 +5,27 @@ module ActiveRecordDataLoader
     class ModelDataGenerator
       attr_reader :table
 
-      def initialize(model:, column_settings:, polymorphic_settings: [], belongs_to_settings: [])
+      def initialize(
+        model:,
+        column_settings:,
+        connection_factory:,
+        logger:,
+        raise_on_duplicates:,
+        max_duplicate_retries:,
+        polymorphic_settings: [],
+        belongs_to_settings: []
+      )
         @model_class = model
         @table = model.table_name
         @column_settings = column_settings
         @polymorphic_settings = polymorphic_settings
         @belongs_to_settings = belongs_to_settings.map { |s| [s.name, s.query] }.to_h
+        @connection_factory = connection_factory
+        @raise_on_duplicates = raise_on_duplicates
+        @max_duplicate_retries = max_duplicate_retries
+        @logger = logger
+        @index_tracker = UniqueIndexTracker.new(model: model, connection_factory: connection_factory)
+        @index_tracker.map_indexed_columns(column_list)
       end
 
       def column_list
@@ -18,11 +33,41 @@ module ActiveRecordDataLoader
       end
 
       def generate_row(row_number)
-        column_list.map { |c| column_data(row_number, c) }
+        @index_tracker.capture_unique_values(generate_row_with_retries(row_number))
       end
 
       private
 
+      def generate_row_with_retries(row_number)
+        retries = 0
+        while @index_tracker.repeating_unique_values?(row = generate_candidate_row(row_number))
+          if (retries += 1) > @max_duplicate_retries
+            raise DuplicateKeyError, <<~MSG if @raise_on_duplicates
+              Exhausted retries looking for unique values for row #{row_number} for '#{table}'.
+              Table '#{table}' has unique indexes that would have prevented inserting this row. If you would
+              like to skip non-unique rows instead of raising, configure `raise_on_duplicates` to be `false`.
+            MSG
+
+            @logger.warn(
+              "[ActiveRecordDataLoader] "\
+              "Exhausted retries looking for unique values. Skipping row #{row_number} for '#{table}'."
+            )
+            return nil
+          else
+            @logger.info(
+              "[ActiveRecordDataLoader] "\
+              "Retrying row #{row_number} for '#{table}' looking for unique values compliant with indexes. "\
+              "Retry number #{retries}."
+            )
+          end
+        end
+        row
+      end
+
+      def generate_candidate_row(row_number)
+        column_list.map { |c| column_data(row_number, c) }
+      end
+
       def column_data(row_number, column)
         column_value = columns[column]
         return column_value unless column_value.respond_to?(:call)
@@ -49,8 +94,14 @@ module ActiveRecordDataLoader
         @model_class
           .columns_hash
           .reject { |name| name == @model_class.primary_key }
-          .select { |_, c| ColumnConfiguration.supported?(model_class: @model_class, ar_column: c) }
-          .map { |_, c| ColumnConfiguration.config_for(model_class: @model_class, ar_column: c) }
+          .select { |_, c| ColumnDataProvider.supported?(model_class: @model_class, ar_column: c) }
+          .map do |_, c|
+            ColumnDataProvider.provider_for(
+              model_class: @model_class,
+              ar_column: c,
+              connection_factory: @connection_factory
+            )
+          end
           .reduce({}, :merge)
       end
 
@@ -60,16 +111,33 @@ module ActiveRecordDataLoader
           .select(&:belongs_to?)
           .reject(&:polymorphic?)
           .map do |assoc|
-            BelongsToConfiguration.config_for(ar_association: assoc, query: @belongs_to_settings[assoc.name])
+            BelongsToDataProvider.provider_for(
+              ar_association: assoc,
+              query: @belongs_to_settings[assoc.name],
+              strategy: column_config_strategy(assoc)
+            )
           end
           .reduce({}, :merge)
       end
 
       def polymorphic_config
         @polymorphic_settings
-          .map { |s| PolymorphicBelongsToConfiguration.config_for(polymorphic_settings: s) }
+          .map do |s|
+            PolymorphicBelongsToDataProvider.provider_for(
+              polymorphic_settings: s,
+              strategy: column_config_strategy(s.model_class.reflect_on_association(s.name))
+            )
+          end
           .reduce({}, :merge)
       end
+
+      def column_config_strategy(column)
+        if @index_tracker.contained_in_index?(column)
+          :cycle
+        else
+          :random
+        end
+      end
     end
   end
 end

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/text_value_generator.rb

@@ -12,7 +12,7 @@ module ActiveRecordDataLoader
       }.freeze
 
       class << self
-        def generator_for(model_class:, ar_column:)
+        def generator_for(model_class:, ar_column:, connection_factory: nil)
           scenario = GENERATORS.keys.find { |m| send(m, model_class, ar_column) }
           generator = GENERATORS.fetch(scenario, -> { SecureRandom.uuid })
 

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,27 +2,52 @@
 
 module ActiveRecordDataLoader
   class Configuration
-    attr_accessor :default_batch_size, :default_row_count, :logger, :statement_timeout
+    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"
+      statement_timeout: "2min",
+      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