active_record_data_loader 0.1.1 → 0.1.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: fbc38165d9ffa4debdc127641a9e79c4fd3eb382dd9271c7e0eed849992c8c72
-  data.tar.gz: d2703bbd36cdbcae163dbf19082704e39dc717499689d9754bc570d846895ada
+  metadata.gz: 43867d09da0479be24e9a23687d53cd4137f5ff8b47f769183a9bee19d6f8372
+  data.tar.gz: 71ea6ea8d1cbc9821168a0f87e90c7e6e08f2b2053c48905bb0bcfe95f5d37f0
 SHA512:
-  metadata.gz: a999728b9926d7e87cb93fbf3329902becc4dd0b51139f38b473820010b93987ecb5e87be0be214111f28bd6c64685af6d2b09d61cc35cf4d3c41dd1fc88f026
-  data.tar.gz: 1b6bcbf2edd88700ce91083ce7da1958ef314f65f5c57932cd978915e36603788696eb3f1cc5894f36dd6870d6d945ca9e72068e2ac05444f5992df099c4760a
+  metadata.gz: 2523c1db84ecd726ba333369646a0fe1e9b6ee8e0e0dc5694e2b8a8ef693fe752c8685914d122c7d9561dd05ff3c1d1905a328445a7407c82596b72cf0f5e9ff
+  data.tar.gz: 3c95553f3b39a535739e3fd998f67ca0c5420e37411d229b9284c0ccc99c9b7690d56a835b8fdaadfddc6f8750604f29852f77198b5a178c4cbab00bf3d80a1f

data/Appraisals

@@ -4,10 +4,6 @@ appraise "activerecord-5" do
   gem "activerecord", "~> 5.0"
 end
 
-appraise "activerecord-6" do
-  gem "activerecord", "6.0.0.rc1"
-end
-
 appraise "faker" do
   gem "faker", ">= 1.9.3"
 end

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    active_record_data_loader (0.1.1)
+    active_record_data_loader (0.1.2)
       activerecord (>= 4.0)
 
 GEM

data/README.md

@@ -1,7 +1,8 @@
 # ActiveRecord Data Loader
 
+[![Gem Version](https://badge.fury.io/rb/active_record_data_loader.svg)](https://badge.fury.io/rb/active_record_data_loader)
 [![Build Status](https://travis-ci.org/abeiderman/active_record_data_loader.svg?branch=master)](https://travis-ci.org/abeiderman/active_record_data_loader)
-[![Coverage Status](https://coveralls.io/repos/github/abeiderman/active_record_data_loader/badge.svg?branch=master)](https://coveralls.io/github/abeiderman/active_record_data_loader?branch=master)
+[![Coverage Status](https://coveralls.io/repos/github/abeiderman/active_record_data_loader/badge.svg?branch=master&service=github)](https://coveralls.io/github/abeiderman/active_record_data_loader?branch=master)
 [![Maintainability](https://api.codeclimate.com/v1/badges/338904b3f7e8d19a3cb1/maintainability)](https://codeclimate.com/github/abeiderman/active_record_data_loader/maintainability)
 
 Efficiently bulk load data for your ActiveRecord models with a simple DSL.

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/lib/active_record_data_loader.rb

@@ -11,6 +11,7 @@ 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"
@@ -45,14 +46,15 @@ module ActiveRecordDataLoader
         generator = ActiveRecordDataLoader::ActiveRecord::ModelDataGenerator.new(
           model: m.klass,
           column_settings: m.columns,
-          polymorphic_settings: m.polymorphic_associations
+          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,
-          logger: configuration.logger
+          configuration: configuration
         )
       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/model_data_generator.rb

@@ -5,11 +5,12 @@ module ActiveRecordDataLoader
     class ModelDataGenerator
       attr_reader :table
 
-      def initialize(model:, column_settings:, polymorphic_settings: [])
+      def initialize(model:, column_settings:, 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
       end
 
       def column_list
@@ -58,7 +59,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/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,18 @@
 
 module ActiveRecordDataLoader
   class Configuration
-    attr_accessor :default_batch_size, :default_row_count, :logger
+    attr_accessor :default_batch_size, :default_row_count, :logger, :statement_timeout
 
     def initialize(
       default_batch_size: 100_000,
       default_row_count: 1,
-      logger: Logger.new(STDOUT, level: :info)
+      logger: Logger.new(STDOUT, level: :info),
+      statement_timeout: "2min"
     )
       @default_batch_size = default_batch_size
       @default_row_count = default_row_count
       @logger = logger
+      @statement_timeout = statement_timeout
     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/loader.rb

@@ -9,10 +9,11 @@ module ActiveRecordDataLoader
         data_generator:,
         total_rows:,
         batch_size:,
-        logger:
+        configuration:
       )
         new(
-          logger: logger,
+          logger: configuration.logger,
+          statement_timeout: configuration.statement_timeout,
           strategy: strategy_class.new(data_generator)
         ).load_data(batch_size, total_rows)
       end
@@ -28,9 +29,10 @@ module ActiveRecordDataLoader
       end
     end
 
-    def initialize(logger:, strategy:)
+    def initialize(logger:, statement_timeout:, strategy:)
       @logger = logger
       @strategy = strategy
+      @statement_timeout = statement_timeout
     end
 
     def load_data(batch_size, total_rows)
@@ -51,16 +53,37 @@ module ActiveRecordDataLoader
 
     private
 
-    attr_reader :strategy, :logger
+    attr_reader :strategy, :statement_timeout, :logger
 
     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(
+            "Completed batch #{i + 1}/#{batch_count}, #{row_numbers.count} row(s) in #{time} seconds"
+          )
+        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 = "0.1.1"
+  VERSION = "0.1.2"
 end

data/script/ci_build.sh

@@ -1,6 +1,4 @@
 #! /bin/bash
 bundle exec appraisal
-bundle exec appraisal activerecord-5 rake
-bundle exec appraisal faker rake
-bundle exec appraisal ffaker rake
-bundle exec rake
+bundle exec appraisal rake
+bundle exec rake coveralls:push

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: active_record_data_loader
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.1.2
 platform: ruby
 authors:
 - Alejandro Beiderman
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2019-06-10 00:00:00.000000000 Z
+date: 2019-06-11 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -206,6 +206,7 @@ files:
 - 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