active_record_data_loader 1.0.0 → 1.2.0

data/lib/active_record_data_loader/active_record/column_configuration.rb

@@ -9,15 +9,17 @@ module ActiveRecordDataLoader
           integer: IntegerValueGenerator,
           string: TextValueGenerator,
           text: TextValueGenerator,
+          datetime: DatetimeValueGenerator,
         }.freeze
 
-        def config_for(model_class:, ar_column:)
+        def config_for(model_class:, ar_column:, connection_factory:)
           raise_error_if_not_supported(model_class, ar_column)
 
           {
-            ar_column.name.to_sym => VALUE_GENERATORS[ar_column.type].generator_for(
+            ar_column.name.to_sym => VALUE_GENERATORS[column_type(ar_column)].generator_for(
               model_class: model_class,
-              ar_column: ar_column
+              ar_column: ar_column,
+              connection_factory: connection_factory
             ),
           }
         end
@@ -25,7 +27,7 @@ module ActiveRecordDataLoader
         def supported?(model_class:, ar_column:)
           return false if model_class.reflect_on_association(ar_column.name)
 
-          VALUE_GENERATORS.keys.include?(ar_column.type)
+          VALUE_GENERATORS.keys.include?(column_type(ar_column))
         end
 
         private
@@ -37,6 +39,14 @@ module ActiveRecordDataLoader
             Column '#{ar_column.name}' of type '#{ar_column.type}' in model '#{model_class.name}' not supported"
           ERROR
         end
+
+        def column_type(ar_column)
+          if ar_column.type == :string && ar_column.sql_type.to_s.downcase.start_with?("enum")
+            :enum
+          else
+            ar_column.type
+          end
+        end
       end
     end
   end

data/lib/active_record_data_loader/active_record/datetime_value_generator.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module ActiveRecordDataLoader
+  module ActiveRecord
+    class DatetimeValueGenerator
+      class << self
+        def generator_for(model_class:, ar_column:, connection_factory: nil)
+          ->(row) { timestamp(model_class, row) }
+        end
+
+        private
+
+        def timestamp(model, row_number)
+          PerRowValueCache[:datetime].get_or_set(model: model, row: row_number) do
+            Time.now.utc
+          end
+        end
+      end
+    end
+  end
+end

data/lib/active_record_data_loader/active_record/enum_value_generator.rb

@@ -4,21 +4,44 @@ module ActiveRecordDataLoader
   module ActiveRecord
     class EnumValueGenerator
       class << self
-        def generator_for(model_class:, ar_column:)
-          values = enum_values_for(model_class, ar_column.sql_type)
+        def generator_for(model_class:, ar_column:, connection_factory:)
+          values = enum_values_for(ar_column.sql_type, connection_factory)
           -> { values.sample }
         end
 
         private
 
-        def enum_values_for(model_class, enum_type)
-          model_class
-            .connection
+        def enum_values_for(enum_type, connection_factory)
+          connection = connection_factory.call
+
+          if connection.adapter_name.downcase.to_sym == :postgresql
+            postgres_enum_values_for(connection, enum_type)
+          elsif connection.adapter_name.downcase.to_s.start_with?("mysql")
+            mysql_enum_values_for(enum_type)
+          else
+            []
+          end
+        ensure
+          connection&.close
+        end
+
+        def postgres_enum_values_for(connection, enum_type)
+          connection
             .execute("SELECT unnest(enum_range(NULL::#{enum_type}))::text")
             .map(&:values)
             .flatten
             .compact
         end
+
+        def mysql_enum_values_for(enum_type)
+          enum_type
+            .to_s
+            .downcase
+            .gsub(/\Aenum\(|\)\Z/, "")
+            .split(",")
+            .map(&:strip)
+            .map { |s| s.gsub(/\A'|'\Z/, "") }
+        end
       end
     end
   end

data/lib/active_record_data_loader/active_record/integer_value_generator.rb

@@ -4,8 +4,8 @@ module ActiveRecordDataLoader
   module ActiveRecord
     class IntegerValueGenerator
       class << self
-        def generator_for(model_class:, ar_column:)
-          range_limit = [(256**number_of_bytes(ar_column)) / 2 - 1, 1_000_000_000].min
+        def generator_for(model_class:, ar_column:, connection_factory: nil)
+          range_limit = [((256**number_of_bytes(ar_column)) / 2) - 1, 1_000_000_000].min
 
           -> { rand(0..range_limit) }
         end

data/lib/active_record_data_loader/active_record/model_data_generator.rb

@@ -5,12 +5,19 @@ module ActiveRecordDataLoader
     class ModelDataGenerator
       attr_reader :table
 
-      def initialize(model:, column_settings:, polymorphic_settings: [], belongs_to_settings: [])
+      def initialize(
+        model:,
+        column_settings:,
+        connection_factory:,
+        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
       end
 
       def column_list
@@ -50,7 +57,13 @@ module ActiveRecordDataLoader
           .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) }
+          .map do |_, c|
+            ColumnConfiguration.config_for(
+              model_class: @model_class,
+              ar_column: c,
+              connection_factory: @connection_factory
+            )
+          end
           .reduce({}, :merge)
       end
 

data/lib/active_record_data_loader/active_record/per_row_value_cache.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module ActiveRecordDataLoader
+  module ActiveRecord
+    class PerRowValueCache
+      class << self
+        def [](key)
+          caches[key] ||= new
+        end
+
+        def clear
+          @caches = {}
+        end
+
+        private
+
+        def caches
+          @caches ||= clear
+        end
+      end
+
+      def initialize
+        @row_caches = Hash.new { |hash, key| hash[key] = {} }
+      end
+
+      def get_or_set(model:, row:)
+        @row_caches[model.name].shift if @row_caches[model.name].size > 1
+
+        @row_caches[model.name][row] ||= yield
+      end
+    end
+  end
+end

data/lib/active_record_data_loader/active_record/polymorphic_belongs_to_configuration.rb

@@ -46,7 +46,7 @@ module ActiveRecordDataLoader
       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/bulk_insert_strategy.rb

@@ -2,12 +2,13 @@
 
 module ActiveRecordDataLoader
   class BulkInsertStrategy
-    def initialize(data_generator)
+    def initialize(data_generator, output_adapter)
       @data_generator = data_generator
+      @output_adapter = output_adapter
     end
 
     def load_batch(row_numbers, connection)
-      connection.insert(<<~SQL)
+      output_adapter.insert(connection: connection, command: <<~SQL)
         INSERT INTO #{quoted_table_name(connection)} (#{column_list(connection)})
         VALUES #{values(row_numbers, connection)}
       SQL
@@ -23,7 +24,7 @@ module ActiveRecordDataLoader
 
     private
 
-    attr_reader :data_generator
+    attr_reader :data_generator, :output_adapter
 
     def quoted_table_name(connection)
       @quoted_table_name ||= connection.quote_table_name(data_generator.table)

data/lib/active_record_data_loader/configuration.rb

@@ -2,27 +2,69 @@
 
 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, :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 },
+      output: :connection
     )
       @default_batch_size = default_batch_size
       @default_row_count = default_row_count
       @logger = logger || default_logger
       @statement_timeout = statement_timeout
+      @connection_factory = connection_factory
+      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
+      )
     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]])
+      else
+        raise "The output configuration parameter must be either a symbol for :connection or :file, "\
+              "or a hash with more detailed output options."
+      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,74 @@
+# frozen_string_literal: true
+
+module ActiveRecordDataLoader
+  class ConnectionHandler
+    def initialize(connection_factory:, statement_timeout:, output_adapter:)
+      @connection_factory = connection_factory
+      @statement_timeout = statement_timeout
+      @output_adapter = output_adapter
+    end
+
+    def with_connection
+      connection = open_connection
+      if postgres?(connection)
+        original_timeout = retrieve_statement_timeout(connection)
+        update_statement_timeout(connection, statement_timeout)
+        yield connection
+        update_statement_timeout(connection, original_timeout)
+      else
+        yield connection
+      end
+    ensure
+      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
+    end
+
+    private
+
+    attr_reader :connection_factory, :statement_timeout, :output_adapter
+
+    def retrieve_statement_timeout(connection)
+      connection.execute("SHOW statement_timeout").first["statement_timeout"]
+    end
+
+    def update_statement_timeout(connection, timeout)
+      connection.execute(statement_timeout_set_command(timeout))
+    end
+
+    def statement_timeout_set_command(timeout)
+      "SET statement_timeout = \"#{timeout}\""
+    end
+
+    def open_connection
+      connection_factory.call
+    end
+
+    def postgres?(connection)
+      connection.adapter_name.downcase.to_sym == :postgresql
+    end
+  end
+end

data/lib/active_record_data_loader/connection_output_adapter.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module ActiveRecordDataLoader
+  class ConnectionOutputAdapter
+    def needs_timeout_output?
+      false
+    end
+
+    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 insert(connection:, command:)
+      connection.insert(command)
+    end
+  end
+end

data/lib/active_record_data_loader/copy_strategy.rb

@@ -2,15 +2,19 @@
 
 module ActiveRecordDataLoader
   class CopyStrategy
-    def initialize(data_generator)
+    def initialize(data_generator, output_adapter)
       @data_generator = data_generator
+      @output_adapter = output_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) }
+      output_adapter.copy(
+        connection: connection,
+        table: table_name_for_copy(connection),
+        columns: columns_for_copy(connection),
+        data: csv_rows(row_numbers, connection),
+        row_numbers: row_numbers
+      )
     end
 
     def table_name
@@ -23,27 +27,23 @@ module ActiveRecordDataLoader
 
     private
 
-    attr_reader :data_generator
+    attr_reader :data_generator, :output_adapter
 
-    def csv_data_batch(row_numbers, connection)
+    def csv_rows(row_numbers, connection)
       row_numbers.map do |i|
         data_generator.generate_row(i).map { |d| quote_data(d, connection) }.join(",")
-      end.join("\n")
+      end
     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(", ")
+    def table_name_for_copy(connection)
+      @table_name_for_copy ||= connection.quote_table_name(data_generator.table)
+    end
 
-        <<~SQL
-          COPY #{quoted_table_name} (#{columns})
-            FROM STDIN WITH (FORMAT CSV)
-        SQL
-      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)

data/lib/active_record_data_loader/file_output_adapter.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module ActiveRecordDataLoader
+  class FileOutputAdapter
+    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))
+    end
+
+    def needs_timeout_output?
+      true
+    end
+
+    def copy(connection:, 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(connection:, command:)
+      execute(command)
+    end
+
+    def execute(command)
+      File.open(@filename, "a") { |f| f.puts("#{command.gsub("\n", ' ')};") }
+    end
+
+    private
+
+    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

@@ -13,15 +13,18 @@ module ActiveRecordDataLoader
       )
         new(
           logger: configuration.logger,
-          statement_timeout: configuration.statement_timeout,
-          strategy: strategy_class.new(data_generator)
+          connection_handler: configuration.connection_handler,
+          strategy: strategy_class(configuration.connection_factory).new(
+            data_generator,
+            configuration.output_adapter
+          )
         ).load_data(batch_size, total_rows)
       end
 
       private
 
-      def strategy_class
-        if ::ActiveRecord::Base.connection.raw_connection.respond_to?(:copy_data)
+      def strategy_class(connection_factory)
+        if connection_factory.call.raw_connection.respond_to?(:copy_data)
           ActiveRecordDataLoader::CopyStrategy
         else
           ActiveRecordDataLoader::BulkInsertStrategy
@@ -29,10 +32,10 @@ module ActiveRecordDataLoader
       end
     end
 
-    def initialize(logger:, statement_timeout:, strategy:)
+    def initialize(logger:, connection_handler:, strategy:)
       @logger = logger
+      @connection_handler = connection_handler
       @strategy = strategy
-      @statement_timeout = statement_timeout
     end
 
     def load_data(batch_size, total_rows)
@@ -55,10 +58,10 @@ module ActiveRecordDataLoader
 
     private
 
-    attr_reader :strategy, :statement_timeout, :logger
+    attr_reader :strategy, :connection_handler, :logger
 
     def load_in_batches(batch_size, total_rows, batch_count)
-      with_connection do |connection|
+      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) }
 
@@ -69,24 +72,5 @@ module ActiveRecordDataLoader
         end
       end
     end
-
-    def with_connection
-      if ::ActiveRecord::Base.connection.adapter_name.downcase.to_sym == :postgresql
-        original_timeout = retrieve_statement_timeout
-        update_statement_timeout(statement_timeout)
-        yield ::ActiveRecord::Base.connection
-        update_statement_timeout(original_timeout)
-      else
-        yield ::ActiveRecord::Base.connection
-      end
-    end
-
-    def retrieve_statement_timeout
-      ::ActiveRecord::Base.connection.execute("SHOW statement_timeout").first["statement_timeout"]
-    end
-
-    def update_statement_timeout(timeout)
-      ::ActiveRecord::Base.connection.execute("SET statement_timeout = \"#{timeout}\"")
-    end
   end
 end

data/lib/active_record_data_loader/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module ActiveRecordDataLoader
-  VERSION = "1.0.0"
+  VERSION = "1.2.0"
 end

data/lib/active_record_data_loader.rb

@@ -3,10 +3,13 @@
 require "active_record_data_loader/version"
 require "active_record"
 require "active_record_data_loader/configuration"
+require "active_record_data_loader/connection_handler"
 require "active_record_data_loader/data_faker"
+require "active_record_data_loader/active_record/per_row_value_cache"
 require "active_record_data_loader/active_record/integer_value_generator"
 require "active_record_data_loader/active_record/text_value_generator"
 require "active_record_data_loader/active_record/enum_value_generator"
+require "active_record_data_loader/active_record/datetime_value_generator"
 require "active_record_data_loader/active_record/column_configuration"
 require "active_record_data_loader/active_record/belongs_to_configuration"
 require "active_record_data_loader/active_record/polymorphic_belongs_to_configuration"
@@ -15,6 +18,8 @@ require "active_record_data_loader/dsl/belongs_to_association"
 require "active_record_data_loader/dsl/polymorphic_association"
 require "active_record_data_loader/dsl/model"
 require "active_record_data_loader/dsl/definition"
+require "active_record_data_loader/connection_output_adapter"
+require "active_record_data_loader/file_output_adapter"
 require "active_record_data_loader/copy_strategy"
 require "active_record_data_loader/bulk_insert_strategy"
 require "active_record_data_loader/loader"
@@ -22,7 +27,7 @@ require "active_record_data_loader/loader"
 module ActiveRecordDataLoader
   def self.define(config = ActiveRecordDataLoader.configuration, &block)
     LoaderProxy.new(
-      configuration,
+      config,
       ActiveRecordDataLoader::Dsl::Definition.new(config).tap { |l| l.instance_eval(&block) }
     )
   end
@@ -42,25 +47,32 @@ module ActiveRecordDataLoader
     end
 
     def load_data
-      definition.models.map do |m|
-        generator = ActiveRecordDataLoader::ActiveRecord::ModelDataGenerator.new(
-          model: m.klass,
-          column_settings: m.columns,
-          polymorphic_settings: m.polymorphic_associations,
-          belongs_to_settings: m.belongs_to_associations
-        )
+      ActiveRecordDataLoader::ActiveRecord::PerRowValueCache.clear
 
-        ActiveRecordDataLoader::Loader.load_data(
-          data_generator: generator,
-          batch_size: m.batch_size,
-          total_rows: m.row_count,
-          configuration: configuration
-        )
+      configuration.connection_handler.with_statement_timeout_for_output do
+        definition.models.map { |m| load_model(m) }
       end
     end
 
     private
 
     attr_reader :definition, :configuration
+
+    def load_model(model)
+      generator = 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
+      )
+
+      ActiveRecordDataLoader::Loader.load_data(
+        data_generator: generator,
+        batch_size: model.batch_size,
+        total_rows: model.row_count,
+        configuration: configuration
+      )
+    end
   end
 end