active_record_data_loader 1.0.1 → 1.0.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
-SHA256:
-  metadata.gz: ddff4ba719de66b138820d4e5c257df064e85259e595156f5f25c54622e9c60a
-  data.tar.gz: e82e4ae901cd557a197c4b5c66069bbb01e2cf64973c4b885a35661f1ecb05ca
+SHA1:
+  metadata.gz: a76332e648af9f10d665eaab3fb442b18d708914
+  data.tar.gz: 012eb43c82d6bca82f97ece0d719faa07ea79c6c
 SHA512:
-  metadata.gz: 6a6d299daed5d81cd92f65206307ceab39e4147afa8956822890941426e4765b197d149d97992d19bc7506697a76cd9e6056623549eeb632cc884dc03f02cea5
-  data.tar.gz: a68aca773f47b9b06f833ad0a8e6eaf6bb29b76d9d966597c17f40855430c799703d75cc8ef34fa209e66cd0cc50b220e00b51b086b1cb136c92607a9f0c873c
+  metadata.gz: 2f41b6aedb4c4cd4f02319cd3d13c50715a0e91bbfd0710e04caa2145054d26b20a4e7a2f69b295b28650a0ad8e111f322fbf008bda28f0db7fa6fa25428de24
+  data.tar.gz: 18d9fb81d1e5d273aab3d21bc2d12f5e28624636f6286deeaf0ccadfdbf20afe629583c96dde78ee32acaaa0b931a363fe32c2c8a77435eca32540875a8d3ee7

data/.travis.yml

@@ -19,5 +19,6 @@ notifications:
 before_install: "gem update --system && gem install bundler"
 before_script:
   - psql -c 'create database test;' -U postgres
+  - mysql -e 'CREATE DATABASE IF NOT EXISTS test;'
   - cp config/database.yml.travis config/database.yml
 script: "bundle exec rake"

data/CHANGELOG.md

@@ -1,5 +1,13 @@
 # Change log
 
+## [v1.0.2] - 2019-07-05
+
+[Diff](https://github.com/abeiderman/active_record_data_loader/compare/v1.0.1...v1.0.2)
+
+### Enhancements:
+* Add support for MySQL enums
+* Accept a connection factory lambda as part of the configuration
+
 ## [v1.0.1] - 2019-06-16
 
 [Diff](https://github.com/abeiderman/active_record_data_loader/compare/v1.0.0...v1.0.1)

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    active_record_data_loader (1.0.1)
+    active_record_data_loader (1.0.2)
       activerecord (>= 4.0)
 
 GEM
@@ -40,6 +40,7 @@ GEM
     json (2.2.0)
     method_source (0.9.2)
     minitest (5.11.3)
+    mysql2 (0.5.2)
     parallel (1.17.0)
     parser (2.6.3.0)
       ast (~> 2.4.0)
@@ -96,6 +97,7 @@ DEPENDENCIES
   appraisal
   bundler (>= 1.16)
   coveralls
+  mysql2
   pg
   pry
   rake (~> 12.0)

data/active_record_data_loader.gemspec

@@ -37,6 +37,7 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "appraisal"
   spec.add_development_dependency "bundler", ">= 1.16"
   spec.add_development_dependency "coveralls"
+  spec.add_development_dependency "mysql2"
   spec.add_development_dependency "pg"
   spec.add_development_dependency "pry"
   spec.add_development_dependency "rake", "~> 12.0"

data/config/database.yml

@@ -8,3 +8,10 @@ postgres:
 sqlite3:
   adapter: "sqlite3"
   database: ":memory:"
+
+mysql:
+  adapter: "mysql2"
+  host: "127.0.0.1"
+  database: "test"
+  username: "test"
+  password: "test"

data/config/database.yml.travis

@@ -5,3 +5,8 @@ postgres:
 sqlite3:
   adapter: "sqlite3"
   database: ":memory:"
+
+mysql:
+  adapter: "mysql2"
+  database: "test"
+  username: "travis"

data/lib/active_record_data_loader.rb

@@ -24,7 +24,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
@@ -46,25 +46,28 @@ module ActiveRecordDataLoader
     def load_data
       ActiveRecordDataLoader::ActiveRecord::PerRowValueCache.clear
 
-      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::Loader.load_data(
-          data_generator: generator,
-          batch_size: m.batch_size,
-          total_rows: m.row_count,
-          configuration: configuration
-        )
-      end
+      definition.models.map { |m| load_model(m) }
     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

data/lib/active_record_data_loader/active_record/column_configuration.rb

@@ -12,13 +12,14 @@ module ActiveRecordDataLoader
           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
@@ -26,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
@@ -38,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

@@ -4,7 +4,7 @@ module ActiveRecordDataLoader
   module ActiveRecord
     class DatetimeValueGenerator
       class << self
-        def generator_for(model_class:, ar_column:)
+        def generator_for(model_class:, ar_column:, connection_factory: nil)
           ->(row) { timestamp(model_class, row) }
         end
 

data/lib/active_record_data_loader/active_record/enum_value_generator.rb

@@ -4,14 +4,26 @@ 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(model_class, ar_column.sql_type, connection_factory)
           -> { values.sample }
         end
 
         private
 
-        def enum_values_for(model_class, enum_type)
+        def enum_values_for(model_class, enum_type, connection_factory)
+          connection = connection_factory.call
+
+          if connection.adapter_name.downcase.to_sym == :postgresql
+            postgres_enum_values_for(model_class, enum_type)
+          elsif connection.adapter_name.downcase.to_s.start_with?("mysql")
+            mysql_enum_values_for(model_class, enum_type)
+          else
+            []
+          end
+        end
+
+        def postgres_enum_values_for(model_class, enum_type)
           model_class
             .connection
             .execute("SELECT unnest(enum_range(NULL::#{enum_type}))::text")
@@ -19,6 +31,16 @@ module ActiveRecordDataLoader
             .flatten
             .compact
         end
+
+        def mysql_enum_values_for(_model_class, 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,7 +4,7 @@ module ActiveRecordDataLoader
   module ActiveRecord
     class IntegerValueGenerator
       class << self
-        def generator_for(model_class:, ar_column:)
+        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) }

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:,
+        polymorphic_settings: [],
+        belongs_to_settings: [],
+        connection_factory:
+      )
         @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/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/configuration.rb

@@ -2,18 +2,20 @@
 
 module ActiveRecordDataLoader
   class Configuration
-    attr_accessor :default_batch_size, :default_row_count, :logger, :statement_timeout
+    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: nil,
-      statement_timeout: "2min"
+      statement_timeout: "2min",
+      connection_factory: -> { ::ActiveRecord::Base.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
     end
 
     private

data/lib/active_record_data_loader/loader.rb

@@ -14,14 +14,15 @@ module ActiveRecordDataLoader
         new(
           logger: configuration.logger,
           statement_timeout: configuration.statement_timeout,
-          strategy: strategy_class.new(data_generator)
+          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
@@ -29,10 +30,11 @@ module ActiveRecordDataLoader
       end
     end
 
-    def initialize(logger:, statement_timeout:, 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)
@@ -55,7 +57,7 @@ module ActiveRecordDataLoader
 
     private
 
-    attr_reader :strategy, :statement_timeout, :logger
+    attr_reader :strategy, :statement_timeout, :logger, :connection_factory
 
     def load_in_batches(batch_size, total_rows, batch_count)
       with_connection do |connection|
@@ -71,22 +73,26 @@ module ActiveRecordDataLoader
     end
 
     def with_connection
-      if ::ActiveRecord::Base.connection.adapter_name.downcase.to_sym == :postgresql
+      if connection.adapter_name.downcase.to_sym == :postgresql
         original_timeout = retrieve_statement_timeout
         update_statement_timeout(statement_timeout)
-        yield ::ActiveRecord::Base.connection
+        yield connection
         update_statement_timeout(original_timeout)
       else
-        yield ::ActiveRecord::Base.connection
+        yield connection
       end
     end
 
     def retrieve_statement_timeout
-      ::ActiveRecord::Base.connection.execute("SHOW statement_timeout").first["statement_timeout"]
+      connection.execute("SHOW statement_timeout").first["statement_timeout"]
     end
 
     def update_statement_timeout(timeout)
-      ::ActiveRecord::Base.connection.execute("SET 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 = "1.0.1"
+  VERSION = "1.0.2"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: active_record_data_loader
 version: !ruby/object:Gem::Version
-  version: 1.0.1
+  version: 1.0.2
 platform: ruby
 authors:
 - Alejandro Beiderman
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2019-06-17 00:00:00.000000000 Z
+date: 2019-07-06 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -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
@@ -251,7 +265,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.0.3
+rubyforge_project: 
+rubygems_version: 2.6.14
 signing_key: 
 specification_version: 4
 summary: A utility to bulk load test data for performance testing.