Add batch support

Author: Ziad Sawalha

README.md

@@ -290,12 +290,12 @@ You are now done defining your importer, let's run it!
 
 ### Import from a file, path or string
 
-You can import from a file, path or just the CSV content. Please note
-that we currently load the entire file in memory. Feel free to
-contribute if you need to support CSV files with millions of lines! :)
+You can import from a file, path or just the CSV content. To avoid loading
+the entire file in memory, enable `use_batches` and optionally provide a
+`batch_size` (default is 1000).
 
 ```ruby
-import = ImportUserCSV.new(file: my_file)
+import = ImportUserCSV.new(file: my_file, use_batches: true, batch_size: 2_500)
 import = ImportUserCSV.new(path: "tmp/new_users.csv")
 import = ImportUserCSV.new(content: "email,name\nbob@example.com,bob")
 ```

lib/csv_importer/csv_reader.rb

@@ -9,6 +9,15 @@ class CSVReader
     attribute :path, String
     attribute :quote_char, String, default: '"'
     attribute :encoding, String, default: 'UTF-8:UTF-8'
+    attribute :batch_size, Integer, default: 1000
+    attribute :use_batches, Boolean, default: false
+
+    def initialize(*args)
+      super
+      @header = nil
+      @rows = nil
+      @csv_enumerator = nil
+    end
 
     def csv_rows
       @csv_rows ||= begin
@@ -25,12 +34,23 @@ def csv_rows
 
     # Returns the header as an Array of Strings
     def header
-      @header ||= csv_rows.first
+      return @header if @header
+
+      if use_batches
+        io = file_or_content_io
+        separator = detect_separator(io.read(1024))
+        io.rewind
+        @header = CSV.parse_line(io.readline, col_sep: separator, quote_char: quote_char)
+      else
+        @header = csv_rows.first
+      end
     end
 
     # Returns the rows as an Array of Arrays of Strings
     def rows
-      @rows ||= csv_rows[1..-1]
+      return @rows ||= csv_rows[1..-1] unless use_batches
+
+      @csv_enumerator ||= create_csv_enumerator
     end
 
     private
@@ -47,13 +67,52 @@ def read_content
       end
     end
 
+    def create_csv_enumerator
+      Enumerator.new do |yielder|
+        io = file_or_content_io
+        separator = detect_separator(io.read(1024))
+        io.rewind
+
+        CSV.new(
+          io,
+          col_sep: separator,
+          quote_char: quote_char,
+          skip_blanks: true,
+          external_encoding: source_encoding
+        ).each_with_index do |row, index|
+          next if index.zero? # Skip header
+          yielder << encode_and_sanitize_row(row)
+        end
+      end
+    end
+
+    def sanitize_content_stream(io)
+      Enumerator.new do |yielder|
+        io.each_line do |line|
+          yielder << sanitize_content(line)
+        end
+      end
+    end
+
+    def file_or_content_io
+      if content.is_a?(String)
+        StringIO.new(content)
+      elsif file
+        file
+      elsif path
+        File.open(path, 'r')
+      else
+        raise Error, "Please provide content, file, or path"
+      end
+    end
+
     def sanitize_content(csv_content)
       csv_content
         .encode(Encoding.find(source_encoding), invalid: :replace, undef: :replace, replace: '') # Remove invalid byte sequences
         .gsub(/\r\r?\n?/, "\n") # Replaces windows line separators with "\n"
     end
 
-    SEPARATORS = [",", ";", "\t"]
+    SEPARATORS = [",", ";", "\t"].freeze
 
     def detect_separator(csv_content)
       SEPARATORS.min_by do |separator|
@@ -87,6 +146,12 @@ def encode_cells(rows)
       end
     end
 
+    def encode_and_sanitize_row(row)
+      row.map do |cell|
+        cell ? cell.encode(target_encoding).strip : ""
+      end
+    end
+
     def source_encoding
       encoding.split(':').first || 'UTF-8'
     end

lib/csv_importer/runner.rb

@@ -12,11 +12,26 @@ def self.call(*args)
     attribute :rows, Array[Row]
     attribute :when_invalid, Symbol
     attribute :after_save_blocks, Array[Proc], default: []
+    attribute :batch_size, Integer, default: 1000
+    attribute :use_batches, Boolean, default: false
 
     attribute :report, Report, default: proc { Report.new }
 
     ImportAborted = Class.new(StandardError)
 
+    def initialize(attributes = {})
+      super
+      @csv_reader = CSVReader.new(
+        content: attributes[:content],
+        file: attributes[:file],
+        path: attributes[:path],
+        quote_char: attributes[:quote_char],
+        encoding: attributes[:encoding],
+        batch_size: batch_size,
+        use_batches: use_batches
+      )
+    end
+
     # Persist the rows' model and return a `Report`
     def call
       if rows.empty?
@@ -42,37 +57,55 @@ def abort_when_invalid?
     end
 
     def persist_rows!
+      if use_batches
+        process_in_batches
+      else
+        process_all_at_once
+      end
+    end
+
+    def process_in_batches
       transaction do
-        rows.each do |row|
-          tags = []
-
-          if row.model.persisted?
-            tags << :update
-          else
-            tags << :create
-          end
-
-          if row.skip?
-            tags << :skip
-          else
-            if row.model.save
-              tags << :success
-            else
-              tags << :failure
-            end
-          end
-
-          add_to_report(row, tags)
-
-          after_save_blocks.each do |block|
-            case block.arity
-            when 0 then block.call
-            when 1 then block.call(row.model)
-            when 2 then block.call(row.model, row.csv_attributes)
-            else
-              raise ArgumentError, "after_save block of arity #{ block.arity } is not supported"
-            end
-          end
+        rows.each_slice(batch_size) do |batch|
+          batch.each { |row| process_row(row) }
+        end
+      end
+    end
+
+    def process_all_at_once
+      transaction do
+        rows.each { |row| process_row(row) }
+      end
+    end
+
+    def process_row(row)
+      tags = []
+
+      if row.model.persisted?
+        tags << :update
+      else
+        tags << :create
+      end
+
+      if row.skip?
+        tags << :skip
+      else
+        if row.model.save
+          tags << :success
+        else
+          tags << :failure
+        end
+      end
+
+      add_to_report(row, tags)
+
+      after_save_blocks.each do |block|
+        case block.arity
+        when 0 then block.call
+        when 1 then block.call(row.model)
+        when 2 then block.call(row.model, row.csv_attributes)
+        else
+          raise ArgumentError, "after_save block of arity #{ block.arity } is not supported"
         end
       end
     end

spec/csv_importer/csv_reader_spec.rb

@@ -1,4 +1,5 @@
 require "spec_helper"
+require 'stringio'
 
 module CSVImporter
   describe CSVReader do
@@ -44,5 +45,119 @@ module CSVImporter
       reader = CSVReader.new(content: "メール,氏名".encode('SJIS'), encoding: 'SJIS:UTF-8')
       expect(reader.header).to eq ["メール", "氏名"]
     end
+
+    context "with batch processing" do
+      let(:csv_content) do
+        "email,first_name,last_name\n" +
+        "john@example.com,John,Doe\n" +
+        "jane@example.com,Jane,Doe\n" +
+        "bob@example.com,Bob,Smith"
+      end
+
+      let(:reader) { CSVReader.new(content: csv_content, use_batches: true, batch_size: 2) }
+
+      it "returns header correctly" do
+        expect(reader.header).to eq ["email", "first_name", "last_name"]
+      end
+
+      it "returns rows as an enumerator" do
+        expect(reader.rows).to be_a(Enumerator)
+      end
+
+      it "yields rows in batches" do
+        rows = reader.rows.to_a
+        expect(rows.size).to eq 3
+        expect(rows[0]).to eq ["john@example.com", "John", "Doe"]
+        expect(rows[1]).to eq ["jane@example.com", "Jane", "Doe"]
+        expect(rows[2]).to eq ["bob@example.com", "Bob", "Smith"]
+      end
+
+      it "batching mechanism works correctly" do
+        large_csv_content = "email,first_name,last_name\n" + ("john@example.com,John,Doe\n" * 1_000)
+        large_reader = CSVReader.new(content: large_csv_content, use_batches: true, batch_size: 100)
+
+        expect(large_reader.header).to eq ["email", "first_name", "last_name"]
+        expect(large_reader.rows).to be_a(Enumerator)
+
+        row_count = 0
+        large_reader.rows.each do |row|
+          expect(row).to eq ["john@example.com", "John", "Doe"]
+          row_count += 1
+          break if row_count >= 100  # We don't need to process all rows for this test
+        end
+
+        expect(row_count).to eq 100
+      end
+
+      it "processes large files in chunks without loading everything into memory" do
+        chunk_size = 1000
+        total_rows = 2000
+        csv_header = "email,first_name,last_name\n"
+        csv_row = "john@example.com,John,Doe\n"
+
+        # Create a StringIO object that simulates a large CSV file
+        csv_content = StringIO.new.tap do |io|
+          io.puts csv_header
+          total_rows.times { io.puts csv_row }
+          io.rewind
+        end
+
+        large_reader = CSVReader.new(file: csv_content, use_batches: true, batch_size: chunk_size)
+
+        expect(large_reader.header).to eq ["email", "first_name", "last_name"]
+        expect(large_reader.rows).to be_a(Enumerator)
+
+        rows_processed = 0
+        max_rows_per_iteration = 0
+
+        large_reader.rows.each_slice(chunk_size) do |batch|
+          rows_in_this_batch = batch.size
+          max_rows_per_iteration = [max_rows_per_iteration, rows_in_this_batch].max
+          rows_processed += rows_in_this_batch
+        end
+
+        expect(rows_processed).to eq total_rows
+        expect(max_rows_per_iteration).to eq chunk_size
+      end
+
+      it "doesn't load the whole file when reading the header" do
+        class StrictIOSimulator < StringIO
+          def initialize(header, content)
+            super(content)
+            @header = header
+            @read_count = 0
+          end
+
+          def gets
+            @read_count += 1
+            raise "Attempted to read beyond header!" if @read_count > 1
+            @header
+          end
+
+          def readline
+            gets
+          end
+
+          def read(bytes = nil)
+            raise "Attempted to read content!" if @read_count > 0
+            super
+          end
+
+          def rewind
+            super
+            @read_count = 0
+          end
+        end
+
+        header = "email,first_name,last_name\n"
+        content = header + ("john@example.com,John,Doe\n" * 100)
+        strict_io = StrictIOSimulator.new(header, content)
+
+        reader = CSVReader.new(file: strict_io, use_batches: true)
+
+        expect { reader.header }.not_to raise_error
+        expect(reader.header).to eq ["email", "first_name", "last_name"]
+      end
+    end
   end
 end