active_record_data_loader 0.1.1 → 1.1.0

data/Rakefile

@@ -3,8 +3,10 @@
 require "bundler/gem_tasks"
 require "rspec/core/rake_task"
 require "rubocop/rake_task"
+require "coveralls/rake/task"
 
 RSpec::Core::RakeTask.new(:spec)
 RuboCop::RakeTask.new(:rubocop)
+Coveralls::RakeTask.new
 
 task default: [:spec, :rubocop]

data/active_record_data_loader.gemspec

@@ -30,18 +30,20 @@ Gem::Specification.new do |spec|
   spec.executables   = spec.files.grep(%r{^bin/}) { |f| File.basename(f) }
   spec.require_paths = ["lib"]
 
-  spec.required_ruby_version = ">= 2.3.0"
+  spec.required_ruby_version = ">= 2.5.0"
 
-  spec.add_dependency "activerecord", ">= 4.0"
+  spec.add_dependency "activerecord", ">= 5.0"
 
   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"
+  spec.add_development_dependency "rake", "~> 13.0"
   spec.add_development_dependency "rspec", "~> 3.0"
   spec.add_development_dependency "rspec-collection_matchers"
   spec.add_development_dependency "rubocop"
   spec.add_development_dependency "sqlite3"
+  spec.add_development_dependency "timecop"
 end

data/config/database.yml

@@ -1,6 +1,7 @@
 postgres:
   adapter: "postgresql"
   host: "127.0.0.1"
+  port: "2345"
   database: "test"
   username: "test"
   password: "test"
@@ -8,3 +9,11 @@ postgres:
 sqlite3:
   adapter: "sqlite3"
   database: ":memory:"
+
+mysql:
+  adapter: "mysql2"
+  host: "127.0.0.1"
+  port: "3306"
+  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/docker-compose.yml

@@ -0,0 +1,18 @@
+version: "3.9"
+services:
+  postgres:
+    image: postgres:11
+    ports:
+      - "2345:5432"
+    environment:
+      - POSTGRES_USER=test
+      - POSTGRES_PASSWORD=test
+  mysql:
+    image: mysql:5
+    ports:
+      - "3306:3306"
+    environment:
+      - MYSQL_ROOT_PASSWORD=test
+      - MYSQL_USER=test
+      - MYSQL_PASSWORD=test
+      - MYSQL_DATABASE=test

data/gemfiles/activerecord_6.gemfile

@@ -2,6 +2,6 @@
 
 source "https://rubygems.org"
 
-gem "activerecord", "6.0.0.rc1"
+gem "activerecord", "~>6.1"
 
 gemspec path: "../"

data/gemfiles/rails.gemfile

@@ -0,0 +1,7 @@
+# This file was generated by Appraisal
+
+source "https://rubygems.org"
+
+gem "rails"
+
+gemspec path: "../"

data/lib/active_record_data_loader.rb

@@ -4,13 +4,16 @@ require "active_record_data_loader/version"
 require "active_record"
 require "active_record_data_loader/configuration"
 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"
 require "active_record_data_loader/active_record/model_data_generator"
+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"
@@ -21,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
@@ -41,24 +44,30 @@ 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
-        )
+      ActiveRecordDataLoader::ActiveRecord::PerRowValueCache.clear
 
-        ActiveRecordDataLoader::Loader.load_data(
-          data_generator: generator,
-          batch_size: m.batch_size,
-          total_rows: m.row_count,
-          logger: configuration.logger
-        )
-      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/belongs_to_configuration.rb

@@ -3,14 +3,15 @@
 module ActiveRecordDataLoader
   module ActiveRecord
     class BelongsToConfiguration
-      def self.config_for(ar_association:)
+      def self.config_for(ar_association:, query: nil)
         raise "#{name} does not support polymorphic associations" if ar_association.polymorphic?
 
-        { ar_association.join_foreign_key.to_sym => new(ar_association).foreign_key_func }
+        { ar_association.join_foreign_key.to_sym => new(ar_association, query).foreign_key_func }
       end
 
-      def initialize(ar_association)
+      def initialize(ar_association, query)
         @ar_association = ar_association
+        @query = query
       end
 
       def foreign_key_func
@@ -20,7 +21,15 @@ module ActiveRecordDataLoader
       private
 
       def possible_values
-        @possible_values ||= @ar_association.klass.all.pluck(@ar_association.join_primary_key).to_a
+        @possible_values ||= base_query.pluck(@ar_association.join_primary_key).to_a
+      end
+
+      def base_query
+        if @query.respond_to?(:call)
+          @query.call.all
+        else
+          @ar_association.klass.all
+        end
       end
     end
   end

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,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,11 +5,19 @@ module ActiveRecordDataLoader
     class ModelDataGenerator
       attr_reader :table
 
-      def initialize(model:, column_settings:, polymorphic_settings: [])
+      def initialize(
+        model:,
+        column_settings:,
+        connection_factory:,
+        polymorphic_settings: [],
+        belongs_to_settings: []
+      )
         @model_class = model
         @table = model.table_name
-        @polymorphic_settings = polymorphic_settings
         @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
@@ -49,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
 
@@ -58,7 +72,9 @@ module ActiveRecordDataLoader
           .reflect_on_all_associations
           .select(&:belongs_to?)
           .reject(&:polymorphic?)
-          .map { |assoc| BelongsToConfiguration.config_for(ar_association: assoc) }
+          .map do |assoc|
+            BelongsToConfiguration.config_for(ar_association: assoc, query: @belongs_to_settings[assoc.name])
+          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

@@ -38,12 +38,20 @@ module ActiveRecordDataLoader
       def possible_values
         @possible_values ||= begin
           values = @settings.models.keys.map do |klass|
-            [klass.name, klass.all.pluck(klass.primary_key).to_a]
+            [klass.name, base_query(klass).pluck(klass.primary_key).to_a]
           end.to_h
 
           @settings.weighted_models.map { |klass| [klass.name, values[klass.name]] }
         end
       end
+
+      def base_query(klass)
+        if @settings.queries[klass].respond_to?(:call)
+          @settings.queries[klass].call.all
+        else
+          klass.all
+        end
+      end
     end
   end
 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 })