active_record_data_loader 0.1.1 → 1.1.0

data/lib/active_record_data_loader/bulk_insert_strategy.rb

@@ -1,17 +1,12 @@
 # frozen_string_literal: true
 
-require "csv"
-require "benchmark"
-
 module ActiveRecordDataLoader
   class BulkInsertStrategy
     def initialize(data_generator)
       @data_generator = data_generator
     end
 
-    def load_batch(row_numbers)
-      connection = ::ActiveRecord::Base.connection
-
+    def load_batch(row_numbers, connection)
       connection.insert(<<~SQL)
         INSERT INTO #{quoted_table_name(connection)} (#{column_list(connection)})
         VALUES #{values(row_numbers, connection)}

data/lib/active_record_data_loader/configuration.rb

@@ -2,16 +2,30 @@
 
 module ActiveRecordDataLoader
   class Configuration
-    attr_accessor :default_batch_size, :default_row_count, :logger
+    attr_accessor :default_batch_size, :default_row_count, :logger, :statement_timeout, :connection_factory
 
     def initialize(
       default_batch_size: 100_000,
       default_row_count: 1,
-      logger: Logger.new(STDOUT, level: :info)
+      logger: nil,
+      statement_timeout: "2min",
+      connection_factory: -> { ::ActiveRecord::Base.connection }
     )
       @default_batch_size = default_batch_size
       @default_row_count = default_row_count
-      @logger = logger
+      @logger = logger || default_logger
+      @statement_timeout = statement_timeout
+      @connection_factory = connection_factory
+    end
+
+    private
+
+    def default_logger
+      if defined?(Rails) && Rails.respond_to?(:logger)
+        Rails.logger
+      else
+        Logger.new($stdout, level: :info)
+      end
     end
   end
 end

data/lib/active_record_data_loader/copy_strategy.rb

@@ -1,19 +1,16 @@
 # frozen_string_literal: true
 
-require "csv"
-require "benchmark"
-
 module ActiveRecordDataLoader
   class CopyStrategy
     def initialize(data_generator)
       @data_generator = data_generator
     end
 
-    def load_batch(row_numbers)
-      csv_data = csv_data_batch(row_numbers)
+    def load_batch(row_numbers, connection)
+      csv_data = csv_data_batch(row_numbers, connection)
 
-      connection = ::ActiveRecord::Base.connection.raw_connection
-      connection.copy_data(copy_command) { connection.put_copy_data(csv_data) }
+      raw_connection = connection.raw_connection
+      raw_connection.copy_data(copy_command(connection)) { raw_connection.put_copy_data(csv_data) }
     end
 
     def table_name
@@ -28,18 +25,18 @@ module ActiveRecordDataLoader
 
     attr_reader :data_generator
 
-    def csv_data_batch(row_numbers)
+    def csv_data_batch(row_numbers, connection)
       row_numbers.map do |i|
-        data_generator.generate_row(i).map { |d| quote_data(d) }.join(",")
+        data_generator.generate_row(i).map { |d| quote_data(d, connection) }.join(",")
       end.join("\n")
     end
 
-    def copy_command
+    def copy_command(connection)
       @copy_command ||= begin
-        quoted_table_name = ::ActiveRecord::Base.connection.quote_table_name(data_generator.table)
+        quoted_table_name = connection.quote_table_name(data_generator.table)
         columns = data_generator
                   .column_list
-                  .map { |c| ::ActiveRecord::Base.connection.quote_column_name(c) }
+                  .map { |c| connection.quote_column_name(c) }
                   .join(", ")
 
         <<~SQL
@@ -49,10 +46,10 @@ module ActiveRecordDataLoader
       end
     end
 
-    def quote_data(data)
+    def quote_data(data, connection)
       return if data.nil?
 
-      "\"#{::ActiveRecord::Base.connection.quote_string(data.to_s)}\""
+      "\"#{connection.quote_string(data.to_s)}\""
     end
   end
 end

data/lib/active_record_data_loader/dsl/belongs_to_association.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module ActiveRecordDataLoader
+  module Dsl
+    class BelongsToAssociation
+      attr_reader :model_class, :name, :query
+
+      def initialize(model_class, name, query = nil)
+        @model_class = model_class
+        @name = name
+        @query = query
+      end
+    end
+  end
+end

data/lib/active_record_data_loader/dsl/model.rb

@@ -3,7 +3,7 @@
 module ActiveRecordDataLoader
   module Dsl
     class Model
-      attr_reader :klass, :columns, :row_count, :polymorphic_associations
+      attr_reader :klass, :columns, :row_count, :polymorphic_associations, :belongs_to_associations
 
       def initialize(klass:, configuration:)
         @klass = klass
@@ -11,6 +11,7 @@ module ActiveRecordDataLoader
         @row_count = configuration.default_row_count
         @batch_size = configuration.default_batch_size
         @polymorphic_associations = []
+        @belongs_to_associations = []
       end
 
       def count(count)
@@ -30,6 +31,10 @@ module ActiveRecordDataLoader
           @klass, assoc_name
         ).tap { |a| block.call(a) }
       end
+
+      def belongs_to(assoc_name, eligible_set:)
+        @belongs_to_associations << BelongsToAssociation.new(@klass, assoc_name, eligible_set)
+      end
     end
   end
 end

data/lib/active_record_data_loader/dsl/polymorphic_association.rb

@@ -3,16 +3,18 @@
 module ActiveRecordDataLoader
   module Dsl
     class PolymorphicAssociation
-      attr_reader :model_class, :name, :models
+      attr_reader :model_class, :name, :models, :queries
 
       def initialize(model_class, name)
         @model_class = model_class
         @name = name
         @models = {}
+        @queries = {}
       end
 
-      def model(klass, weight: 1)
+      def model(klass, weight: 1, eligible_set: nil)
         @models[klass] = weight.to_i
+        @queries[klass] = eligible_set if eligible_set
       end
 
       def weighted_models

data/lib/active_record_data_loader/loader.rb

@@ -9,18 +9,20 @@ module ActiveRecordDataLoader
         data_generator:,
         total_rows:,
         batch_size:,
-        logger:
+        configuration:
       )
         new(
-          logger: logger,
-          strategy: strategy_class.new(data_generator)
+          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
-        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
@@ -28,15 +30,18 @@ module ActiveRecordDataLoader
       end
     end
 
-    def initialize(logger:, strategy:)
+    def initialize(logger:, statement_timeout:, strategy:, connection_factory:)
       @logger = logger
       @strategy = strategy
+      @statement_timeout = statement_timeout
+      @connection_factory = connection_factory
     end
 
     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)."
       )
@@ -44,6 +49,7 @@ module ActiveRecordDataLoader
         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."
       )
@@ -51,16 +57,42 @@ module ActiveRecordDataLoader
 
     private
 
-    attr_reader :strategy, :logger
+    attr_reader :strategy, :statement_timeout, :logger, :connection_factory
 
     def load_in_batches(batch_size, total_rows, batch_count)
-      total_rows.times.each_slice(batch_size).with_index do |row_numbers, i|
-        time = Benchmark.realtime { strategy.load_batch(row_numbers) }
+      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) }
 
-        logger.debug(
-          "Completed batch #{i + 1}/#{batch_count}, #{row_numbers.count} row(s) in #{time} seconds"
-        )
+          logger.debug(
+            "[ActiveRecordDataLoader] "\
+            "Completed batch #{i + 1}/#{batch_count}, #{row_numbers.count} row(s) in #{time} seconds"
+          )
+        end
       end
     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)
+      else
+        yield connection
+      end
+    end
+
+    def retrieve_statement_timeout
+      connection.execute("SHOW statement_timeout").first["statement_timeout"]
+    end
+
+    def update_statement_timeout(timeout)
+      connection.execute("SET statement_timeout = \"#{timeout}\"")
+    end
+
+    def connection
+      connection_factory.call
+    end
   end
 end

data/lib/active_record_data_loader/version.rb

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

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: active_record_data_loader
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 1.1.0
 platform: ruby
 authors:
 - Alejandro Beiderman
-autorequire: 
+autorequire:
 bindir: bin
 cert_chain: []
-date: 2019-06-10 00:00:00.000000000 Z
+date: 2021-05-01 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -16,14 +16,14 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '4.0'
+        version: '5.0'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '4.0'
+        version: '5.0'
 - !ruby/object:Gem::Dependency
   name: appraisal
   requirement: !ruby/object:Gem::Requirement
@@ -66,6 +66,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: mysql2
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
 - !ruby/object:Gem::Dependency
   name: pg
   requirement: !ruby/object:Gem::Requirement
@@ -100,14 +114,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '12.0'
+        version: '13.0'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '12.0'
+        version: '13.0'
 - !ruby/object:Gem::Dependency
   name: rspec
   requirement: !ruby/object:Gem::Requirement
@@ -164,6 +178,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: timecop
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
 description: A utility to bulk load test data for performance testing.
 email:
 - abeiderman@gmail.com
@@ -178,6 +206,7 @@ files:
 - ".rubocop.yml"
 - ".travis.yml"
 - Appraisals
+- CHANGELOG.md
 - CODE_OF_CONDUCT.md
 - Gemfile
 - Gemfile.lock
@@ -189,36 +218,40 @@ files:
 - bin/setup
 - config/database.yml
 - config/database.yml.travis
+- docker-compose.yml
 - gemfiles/.bundle/config
 - gemfiles/activerecord_5.gemfile
 - gemfiles/activerecord_6.gemfile
 - gemfiles/faker.gemfile
 - gemfiles/ffaker.gemfile
+- gemfiles/rails.gemfile
 - lib/active_record_data_loader.rb
 - lib/active_record_data_loader/active_record/belongs_to_configuration.rb
 - lib/active_record_data_loader/active_record/column_configuration.rb
+- lib/active_record_data_loader/active_record/datetime_value_generator.rb
 - lib/active_record_data_loader/active_record/enum_value_generator.rb
 - lib/active_record_data_loader/active_record/integer_value_generator.rb
 - lib/active_record_data_loader/active_record/model_data_generator.rb
+- lib/active_record_data_loader/active_record/per_row_value_cache.rb
 - lib/active_record_data_loader/active_record/polymorphic_belongs_to_configuration.rb
 - lib/active_record_data_loader/active_record/text_value_generator.rb
 - lib/active_record_data_loader/bulk_insert_strategy.rb
 - lib/active_record_data_loader/configuration.rb
 - lib/active_record_data_loader/copy_strategy.rb
 - lib/active_record_data_loader/data_faker.rb
+- lib/active_record_data_loader/dsl/belongs_to_association.rb
 - lib/active_record_data_loader/dsl/definition.rb
 - lib/active_record_data_loader/dsl/model.rb
 - lib/active_record_data_loader/dsl/polymorphic_association.rb
 - lib/active_record_data_loader/loader.rb
 - lib/active_record_data_loader/version.rb
 - log/.keep
-- script/ci_build.sh
-homepage: 
+homepage:
 licenses:
 - MIT
 metadata:
   source_code_uri: https://github.com/abeiderman/active_record_data_loader
-post_install_message: 
+post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -226,15 +259,15 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: 2.3.0
+      version: 2.5.0
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.0.3
-signing_key: 
+rubygems_version: 3.1.4
+signing_key:
 specification_version: 4
 summary: A utility to bulk load test data for performance testing.
 test_files: []