active_record_data_loader 1.1.0 → 1.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 33d3769ad2f8c008cf51e75a62b43ba3333dce8e7e45387e01ab3470f2a5bf6c
-  data.tar.gz: 6f2b812d1e3837d05ff565c050cd6c80110929136601fb95436996b14de4e809
+  metadata.gz: 8ae60492b342e71d8a1a2e20aff31794af784f1009346f4d5dc4375dd939bb48
+  data.tar.gz: ab5be5bdf1cb5af69f5f8a1ac65ea8f64ee1338398c7f6cae6e15e6276d2e954
 SHA512:
-  metadata.gz: 8813fa219e7a710f147b79bb3a81826a5fe23ed7e3bef21a95fca310b3c8f3b09c7fb2c89d0cb3a65c825d01941eca8e394ee42900e13941d050e05fcb2c76d5
-  data.tar.gz: 43392902671095ffcf3d9f33984b7f77de3a9b2fb95bd99ee65378d046e2d509d34f8b92220de04642e6e33d816926215581caa58d02b0d30ac02011f1956776
+  metadata.gz: 8b860a06433a0661a765428f37df596ac0b3fd46ff447fbf3edfa5cf50f89dcda0f146b25074959ca5aee792dff153a56eb37888075ae10649d6cde1e598c53d
+  data.tar.gz: dc9b43011f8c1b8c7b663b650b5ba11e43b9dd040a8ef6b7a58b24c7da903b1c963104a06c39f45e5ebaa4f01e815d8d6d565ded2fd9b5a186dcd50ed69b2e7c

data/.github/workflows/build.yml

@@ -0,0 +1,51 @@
+name: Build
+
+on: [push, workflow_dispatch]
+
+jobs:
+  build:
+    name: Build + Test
+    runs-on: ubuntu-latest
+    services:
+      postgres:
+        image: postgres:11
+        ports:
+          - "2345:5432"
+        env:
+          POSTGRES_USER: test
+          POSTGRES_PASSWORD: test
+      mysql:
+        image: mysql:5
+        ports:
+          - "3306:3306"
+        env:
+          MYSQL_ROOT_PASSWORD: test
+          MYSQL_USER: test
+          MYSQL_PASSWORD: test
+          MYSQL_DATABASE: test
+    strategy:
+      matrix:
+        ruby: [2.5.9, 2.6.7, 2.7.3]
+        gemfile: [activerecord_5, rails, faker, ffaker]
+    env:
+      BUNDLE_GEMFILE: ${{ github.workspace }}/gemfiles/${{ matrix.gemfile }}.gemfile
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v2
+
+      - name: Setup ruby
+        uses: ruby/setup-ruby@v1
+        with:
+          ruby-version: ${{ matrix.ruby }}
+          bundler-cache: true
+
+      - name: Wait for DBs to be ready
+        run: bundle exec rake wait_for_test_db
+
+      - name: Run tests
+        run: bundle exec rake
+
+      - name: Coveralls
+        uses: coverallsapp/github-action@master
+        with:
+          github-token: ${{ secrets.github_token }}

data/.github/workflows/gem-push.yml

@@ -0,0 +1,29 @@
+name: Ruby Gem
+
+on: workflow_dispatch
+
+jobs:
+  build:
+    name: Build + Publish
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Ruby 2.6
+      uses: actions/setup-ruby@v1
+      with:
+        ruby-version: 2.6.x
+
+    - name: Publish to RubyGems
+      run: |
+        mkdir -p $HOME/.gem
+        touch $HOME/.gem/credentials
+        chmod 0600 $HOME/.gem/credentials
+        printf -- "---\n:rubygems_api_key: ${GEM_HOST_API_KEY}\n" > $HOME/.gem/credentials
+        gem build *.gemspec
+        gem push *.gem
+      env:
+        GEM_HOST_API_KEY: "${{ secrets.RUBYGEMS_API_KEY }}"

data/.rubocop.yml

@@ -26,10 +26,17 @@ Metrics/AbcSize:
 Metrics/BlockLength:
   Exclude: ["spec/**/*", "*.gemspec"]
 
+Metrics/ClassLength:
+  Exclude: ["spec/**/*", "*.gemspec"]
+
 Metrics/MethodLength:
   Max: 15
   Exclude: ["spec/**/*"]
 
+Metrics/ParameterLists:
+  Max: 5
+  Exclude: ["lib/active_record_data_loader/configuration.rb"]
+
 Style/CaseLikeIf:
   Enabled: false
 

data/CHANGELOG.md

@@ -1,10 +1,26 @@
 # Change log
 
+## [v1.2.0] - 2021-11-14
+
+[Diff](https://github.com/abeiderman/active_record_data_loader/compare/v1.1.0...v1.2.0)
+
+### Changes:
+* Add `:file` output option for generating a SQL script instead of loading the data into the database.
+* Fix some connection handling issues when a custom connection factory is provided.
+
+## [v1.1.0] - 2021-05-01
+
+[Diff](https://github.com/abeiderman/active_record_data_loader/compare/v1.0.2...v1.1.0)
+
+### Changes:
+* Bump ruby version requirement to >= 2.5
+* Bump activerecord requirement to >= 5.0
+
 ## [v1.0.2] - 2019-07-05
 
 [Diff](https://github.com/abeiderman/active_record_data_loader/compare/v1.0.1...v1.0.2)
 
-### Enhancements:
+### Changes:
 * Add support for MySQL enums
 * Accept a connection factory lambda as part of the configuration
 
@@ -12,7 +28,7 @@
 
 [Diff](https://github.com/abeiderman/active_record_data_loader/compare/v1.0.0...v1.0.1)
 
-### Enhancements:
+### Changes:
 * Generate values for datetime column types. This also fixes the fact that `created_at` and `updated_at` were not being populated by default.
 
 ## [v1.0.0] - 2019-06-15
@@ -22,3 +38,5 @@ Initial stable release
 [v1.0.0]: https://github.com/abeiderman/active_record_data_loader/releases/tag/v1.0.0
 [v1.0.1]: https://github.com/abeiderman/active_record_data_loader/releases/tag/v1.0.1
 [v1.0.2]: https://github.com/abeiderman/active_record_data_loader/releases/tag/v1.0.2
+[v1.1.0]: https://github.com/abeiderman/active_record_data_loader/releases/tag/v1.1.0
+[v1.2.0]: https://github.com/abeiderman/active_record_data_loader/releases/tag/v1.2.0

data/Gemfile.lock

@@ -1,47 +1,39 @@
 PATH
   remote: .
   specs:
-    active_record_data_loader (1.1.0)
+    active_record_data_loader (1.2.0)
       activerecord (>= 5.0)
 
 GEM
   remote: https://rubygems.org/
   specs:
-    activemodel (5.2.4.5)
-      activesupport (= 5.2.4.5)
-    activerecord (5.2.4.5)
-      activemodel (= 5.2.4.5)
-      activesupport (= 5.2.4.5)
-      arel (>= 9.0)
-    activesupport (5.2.4.5)
+    activemodel (6.1.4.1)
+      activesupport (= 6.1.4.1)
+    activerecord (6.1.4.1)
+      activemodel (= 6.1.4.1)
+      activesupport (= 6.1.4.1)
+    activesupport (6.1.4.1)
       concurrent-ruby (~> 1.0, >= 1.0.2)
-      i18n (>= 0.7, < 2)
-      minitest (~> 5.1)
-      tzinfo (~> 1.1)
+      i18n (>= 1.6, < 2)
+      minitest (>= 5.1)
+      tzinfo (~> 2.0)
+      zeitwerk (~> 2.3)
     appraisal (2.2.0)
       bundler
       rake
       thor (>= 0.14.0)
-    arel (9.0.0)
-    ast (2.4.1)
+    ast (2.4.2)
     coderay (1.1.2)
-    concurrent-ruby (1.1.8)
-    coveralls (0.8.23)
-      json (>= 1.8, < 3)
-      simplecov (~> 0.16.1)
-      term-ansicolor (~> 1.3)
-      thor (>= 0.19.4, < 2.0)
-      tins (~> 1.6)
+    concurrent-ruby (1.1.9)
     diff-lcs (1.3)
-    docile (1.3.2)
-    i18n (1.8.10)
+    docile (1.4.0)
+    i18n (1.8.11)
       concurrent-ruby (~> 1.0)
-    json (2.3.1)
     method_source (1.0.0)
     minitest (5.14.4)
     mysql2 (0.5.3)
-    parallel (1.20.1)
-    parser (2.7.2.0)
+    parallel (1.21.0)
+    parser (3.0.2.0)
       ast (~> 2.4.1)
     pg (1.2.3)
     pry (0.13.1)
@@ -49,7 +41,7 @@ GEM
       method_source (~> 1.0)
     rainbow (3.0.0)
     rake (13.0.1)
-    regexp_parser (2.0.0)
+    regexp_parser (2.1.1)
     rexml (3.2.5)
     rspec (3.9.0)
       rspec-core (~> 3.9.0)
@@ -66,33 +58,32 @@ GEM
       diff-lcs (>= 1.2.0, < 2.0)
       rspec-support (~> 3.9.0)
     rspec-support (3.9.3)
-    rubocop (1.6.1)
+    rubocop (1.22.3)
       parallel (~> 1.10)
-      parser (>= 2.7.1.5)
+      parser (>= 3.0.0.0)
       rainbow (>= 2.2.2, < 4.0)
       regexp_parser (>= 1.8, < 3.0)
       rexml
-      rubocop-ast (>= 1.2.0, < 2.0)
+      rubocop-ast (>= 1.12.0, < 2.0)
       ruby-progressbar (~> 1.7)
-      unicode-display_width (>= 1.4.0, < 2.0)
-    rubocop-ast (1.3.0)
-      parser (>= 2.7.1.5)
-    ruby-progressbar (1.10.1)
-    simplecov (0.16.1)
+      unicode-display_width (>= 1.4.0, < 3.0)
+    rubocop-ast (1.13.0)
+      parser (>= 3.0.1.1)
+    ruby-progressbar (1.11.0)
+    simplecov (0.21.2)
       docile (~> 1.1)
-      json (>= 1.8, < 3)
-      simplecov-html (~> 0.10.0)
-    simplecov-html (0.10.2)
+      simplecov-html (~> 0.11)
+      simplecov_json_formatter (~> 0.1)
+    simplecov-html (0.12.3)
+    simplecov-lcov (0.8.0)
+    simplecov_json_formatter (0.1.3)
     sqlite3 (1.4.1)
-    term-ansicolor (1.7.1)
-      tins (~> 1.0)
     thor (0.20.3)
-    thread_safe (0.3.6)
     timecop (0.9.1)
-    tins (1.21.0)
-    tzinfo (1.2.9)
-      thread_safe (~> 0.1)
-    unicode-display_width (1.7.0)
+    tzinfo (2.0.4)
+      concurrent-ruby (~> 1.0)
+    unicode-display_width (2.1.0)
+    zeitwerk (2.5.1)
 
 PLATFORMS
   ruby
@@ -102,7 +93,6 @@ DEPENDENCIES
   active_record_data_loader!
   appraisal
   bundler (>= 1.16)
-  coveralls
   mysql2
   pg
   pry
@@ -110,8 +100,10 @@ DEPENDENCIES
   rspec (~> 3.0)
   rspec-collection_matchers
   rubocop
+  simplecov
+  simplecov-lcov
   sqlite3
   timecop
 
 BUNDLED WITH
-   2.2.1
+   2.2.31

data/README.md

@@ -1,6 +1,6 @@
 # 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)
+[![Build Status](https://github.com/abeiderman/active_record_data_loader/actions/workflows/build.yml/badge.svg)](https://github.com/abeiderman/active_record_data_loader/actions/workflows/build.yml)
 [![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)
 
@@ -99,13 +99,13 @@ data_loader = ActiveRecordDataLoader.define do
     m.column :currency, "CAD"
     m.belongs_to :customer, eligible_set: -> { Customer.where(country: "CAN") }
   end
-  
+
   model Order do |m|
     m.count 25_000
     m.column :currency, "MXN"
     m.belongs_to :customer, eligible_set: -> { Customer.where(country: "MEX") }
   end
-  
+
    model Order do |m|
     m.count 50_000
     m.column :currency, "USD"
@@ -148,7 +148,7 @@ data_loader = ActiveRecordDataLoader.define do
 
   model Order do |m|
     m.count 100_000
-    
+
     m.polymorphic :customer do |c|
       c.model Person
       c.model Business
@@ -172,7 +172,7 @@ data_loader = ActiveRecordDataLoader.define do
 
   model Order do |m|
     m.count 100_000
-    
+
     m.polymorphic :customer do |c|
       c.model Person, weight: 2
       c.model Business, weight: 1
@@ -197,7 +197,7 @@ data_loader = ActiveRecordDataLoader.define do
 
   model Order do |m|
     m.count 100_000
-    
+
     m.polymorphic :customer do |c|
       c.model Person, weight: 2
       c.model Business, weight: 1, eligible_set: -> { Business.where(country: "USA") }
@@ -208,6 +208,85 @@ end
 data_loader.load_data
 ```
 
+### Configuration options
+
+You can define global configuration options like this:
+
+```ruby
+ActiveRecordDataLoader.configure do |c|
+  c.logger = ActiveSupport::Logger.new("my_file.log", level: :debug)
+  c.statement_timeout = "5min"
+end
+```
+
+Or you can create a configuration object for the specific data loader instance rather than globally:
+
+```ruby
+config = ActiveRecordDataLoader::Configuration.new(
+  c.logger = ActiveSupport::Logger.new("my_file.log", level: :debug)
+  c.statement_timeout = "5min"
+)
+loader = ActiveRecordDataLoader.define(config) do
+  model Company do |m|
+    m.count 10
+  end
+
+  # ... more definitions
+end
+```
+
+#### statement_timeout
+
+This is currently only used for Postgres connections to adjust the `statement_timeout` value for the connection. The default is `2min`. Depending on the size of the batches you are loading and overall size of the tables you may need to increase this value:
+
+```ruby
+ActiveRecordDataLoader.configure do |c|
+  c.statement_timeout = "5min"
+end
+```
+
+#### connection_factory
+
+The `connection_factory` option accepts a lambda that should return a connection object whenever executed. If not specified, the default behavior is to retrieve a connection using `ActiveRecord::Base.connection`. You can configure it like this:
+
+```ruby
+ActiveRecordDataLoader.configure do |c|
+  c.connection_factory = -> { MyCustomConnectionHandler.open_connection }
+end
+```
+
+#### output
+
+The `output` option accepts either `:connection` or `:file`. The default behavior is `:connection` which means the data will be loaded into the database using the database connection.
+
+If `:file` is specified, instead of the data being loaded into the database, a script file will be generated. This script file can then be executed manually to load the data. This can be helpful if you need to load the same data multiple times. For example if you are profiling different alternatives in your code and you want to see how each performs with a fully loaded database. In that case you would want to have the same data starting point for each alternative you evaluate. By generating the script file ahead of time, it would be significantly faster to load that data over and over by executing the existing script.
+
+Here are some examples on how to use the `output` option:
+
+```ruby
+ActiveRecordDataLoader.configure do |c|
+  c.output = :connection  # This is the default behavior
+end
+```
+
+```ruby
+ActiveRecordDataLoader.configure do |c|
+  c.output = :file  # Outputs to a file with a default name
+end
+```
+
+```ruby
+ActiveRecordDataLoader.configure do |c|
+  c.output = { type: :file, filename: "./my_script.sql" }  # Outputs to the provided file
+end
+```
+
+When using the `:file` type with Postgres, the resulting script will have `\COPY` commands which reference CSV files that contain the data batches to be copied. The CSV files will be created along side the SQL script and will have a naming convention of using the table name and the rows range for the given batch. For example `./my_script_customers_1_to_1000.csv`. Each `\COPY` command in the SQL file will reference the corresponding CSV file so all you need to do is execute the SQL file using `psql`:
+
+```bash
+psql -h my-db-host -U my_user -f my_script.sql
+```
+
 ## Development
 
 After checking out the repo, run `bin/setup` to install dependencies. Then, run `rake spec` to run the tests. You can also run `bin/console` for an interactive prompt that will allow you to experiment.

data/Rakefile

@@ -3,10 +3,16 @@
 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]
+
+task :wait_for_test_db do
+  require "active_record_data_loader"
+  require "./spec/active_record_helper"
+
+  ActiveRecordHelper.wait_for_mysql
+  ActiveRecordHelper.wait_for_postgres
+end

data/active_record_data_loader.gemspec

@@ -20,7 +20,7 @@ Gem::Specification.new do |spec|
     spec.metadata["source_code_uri"] = "https://github.com/abeiderman/active_record_data_loader"
   else
     raise "RubyGems 2.0 or newer is required to protect against " \
-      "public gem pushes."
+          "public gem pushes."
   end
 
   spec.files = `git ls-files -z`.split("\x0").reject do |f|
@@ -36,7 +36,6 @@ 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"
@@ -44,6 +43,8 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "rspec", "~> 3.0"
   spec.add_development_dependency "rspec-collection_matchers"
   spec.add_development_dependency "rubocop"
+  spec.add_development_dependency "simplecov"
+  spec.add_development_dependency "simplecov-lcov"
   spec.add_development_dependency "sqlite3"
   spec.add_development_dependency "timecop"
 end

data/lib/active_record_data_loader/active_record/enum_value_generator.rb

@@ -5,34 +5,35 @@ module ActiveRecordDataLoader
     class EnumValueGenerator
       class << self
         def generator_for(model_class:, ar_column:, connection_factory:)
-          values = enum_values_for(model_class, ar_column.sql_type, connection_factory)
+          values = enum_values_for(ar_column.sql_type, connection_factory)
           -> { values.sample }
         end
 
         private
 
-        def enum_values_for(model_class, enum_type, connection_factory)
+        def enum_values_for(enum_type, connection_factory)
           connection = connection_factory.call
 
           if connection.adapter_name.downcase.to_sym == :postgresql
-            postgres_enum_values_for(model_class, enum_type)
+            postgres_enum_values_for(connection, enum_type)
           elsif connection.adapter_name.downcase.to_s.start_with?("mysql")
-            mysql_enum_values_for(model_class, enum_type)
+            mysql_enum_values_for(enum_type)
           else
             []
           end
+        ensure
+          connection&.close
         end
 
-        def postgres_enum_values_for(model_class, enum_type)
-          model_class
-            .connection
+        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(_model_class, enum_type)
+        def mysql_enum_values_for(enum_type)
           enum_type
             .to_s
             .downcase

data/lib/active_record_data_loader/active_record/integer_value_generator.rb

@@ -5,7 +5,7 @@ module ActiveRecordDataLoader
     class IntegerValueGenerator
       class << self
         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
+          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/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,24 +2,64 @@
 
 module ActiveRecordDataLoader
   class Configuration
-    attr_accessor :default_batch_size, :default_row_count, :logger, :statement_timeout, :connection_factory
+    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",
-      connection_factory: -> { ::ActiveRecord::Base.connection }
+      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

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,9 +13,11 @@ module ActiveRecordDataLoader
       )
         new(
           logger: configuration.logger,
-          statement_timeout: configuration.statement_timeout,
-          strategy: strategy_class(configuration.connection_factory).new(data_generator),
-          connection_factory: configuration.connection_factory
+          connection_handler: configuration.connection_handler,
+          strategy: strategy_class(configuration.connection_factory).new(
+            data_generator,
+            configuration.output_adapter
+          )
         ).load_data(batch_size, total_rows)
       end
 
@@ -30,11 +32,10 @@ module ActiveRecordDataLoader
       end
     end
 
-    def initialize(logger:, statement_timeout:, strategy:, connection_factory:)
+    def initialize(logger:, connection_handler:, strategy:)
       @logger = logger
+      @connection_handler = connection_handler
       @strategy = strategy
-      @statement_timeout = statement_timeout
-      @connection_factory = connection_factory
     end
 
     def load_data(batch_size, total_rows)
@@ -57,10 +58,10 @@ module ActiveRecordDataLoader
 
     private
 
-    attr_reader :strategy, :statement_timeout, :logger, :connection_factory
+    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) }
 
@@ -71,28 +72,5 @@ module ActiveRecordDataLoader
         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 = "1.1.0"
+  VERSION = "1.2.0"
 end

data/lib/active_record_data_loader.rb

@@ -3,6 +3,7 @@
 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"
@@ -17,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"
@@ -46,7 +49,9 @@ module ActiveRecordDataLoader
     def load_data
       ActiveRecordDataLoader::ActiveRecord::PerRowValueCache.clear
 
-      definition.models.map { |m| load_model(m) }
+      configuration.connection_handler.with_statement_timeout_for_output do
+        definition.models.map { |m| load_model(m) }
+      end
     end
 
     private

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: active_record_data_loader
 version: !ruby/object:Gem::Version
-  version: 1.1.0
+  version: 1.2.0
 platform: ruby
 authors:
 - Alejandro Beiderman
-autorequire:
+autorequire: 
 bindir: bin
 cert_chain: []
-date: 2021-05-01 00:00:00.000000000 Z
+date: 2021-11-14 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -52,20 +52,6 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '1.16'
-- !ruby/object:Gem::Dependency
-  name: coveralls
-  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: mysql2
   requirement: !ruby/object:Gem::Requirement
@@ -164,6 +150,34 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: simplecov
+  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: simplecov-lcov
+  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: sqlite3
   requirement: !ruby/object:Gem::Requirement
@@ -201,10 +215,11 @@ executables:
 extensions: []
 extra_rdoc_files: []
 files:
+- ".github/workflows/build.yml"
+- ".github/workflows/gem-push.yml"
 - ".gitignore"
 - ".rspec"
 - ".rubocop.yml"
-- ".travis.yml"
 - Appraisals
 - CHANGELOG.md
 - CODE_OF_CONDUCT.md
@@ -217,7 +232,6 @@ files:
 - bin/console
 - bin/setup
 - config/database.yml
-- config/database.yml.travis
 - docker-compose.yml
 - gemfiles/.bundle/config
 - gemfiles/activerecord_5.gemfile
@@ -237,21 +251,24 @@ files:
 - 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/connection_handler.rb
+- lib/active_record_data_loader/connection_output_adapter.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/file_output_adapter.rb
 - lib/active_record_data_loader/loader.rb
 - lib/active_record_data_loader/version.rb
 - log/.keep
-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
@@ -266,8 +283,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.1.4
-signing_key:
+rubygems_version: 3.0.3.1
+signing_key: 
 specification_version: 4
 summary: A utility to bulk load test data for performance testing.
 test_files: []

data/.travis.yml

@@ -1,24 +0,0 @@
-sudo: false
-language: ruby
-env:
-  - COVERALLS_PARALLEL=true
-rvm:
-  - 2.5.9
-  - 2.6.7
-  - 2.7.3
-gemfile:
-  - gemfiles/activerecord_5.gemfile
-  - gemfiles/rails.gemfile
-  - gemfiles/faker.gemfile
-  - gemfiles/ffaker.gemfile
-services:
-  - postgresql
-  - mysql
-notifications:
-  webhooks: https://coveralls.io/webhook
-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/config/database.yml.travis

@@ -1,12 +0,0 @@
-postgres:
-  adapter: "postgresql"
-  database: "test"
-
-sqlite3:
-  adapter: "sqlite3"
-  database: ":memory:"
-
-mysql:
-  adapter: "mysql2"
-  database: "test"
-  username: "travis"