tire bulk importer

tire_bulk_import.rake

module Tire
  module Tasks
    module Import
      def bulk_import_model(index, klass, options = {})
        unless progress_bar(klass)
          puts "[IMPORT] Importing '#{klass.to_s}'"
        end
        klass.scoped.find_in_batches(options) do |batch|
          documents = batch.map(&:to_indexed_json)
          progress_bar(klass).inc documents.size if progress_bar(klass)
          index.bulk :index, documents, refresh: true
        end
        progress_bar(klass).finish if progress_bar(klass)
      end
    end
  end
end


namespace :tire_bulk_import do
  desc "Generate api docs"
  task :model do
    klass  = eval(ENV['CLASS'].to_s)
    total  = klass.count rescue nil

    if ENV['INDEX']
      index = Tire::Index.new(ENV['INDEX'])
    else
      index = klass.tire.index
    end

    batch_size = ENV['BATCH_SIZE'].to_i || 1000
    includes = klass.respond_to?(:includes_for_index) ? klass.includes_for_index : []
    options  = {batch_size: batch_size, include: includes}

    Tire::Tasks::Import.progress_bar(klass, total) if total
    Tire::Tasks::Import.delete_index(index) if ENV['FORCE']
    Tire::Tasks::Import.create_index(index, klass)
    puts "[IMPORT] Import Records:   #{total}"
    Tire::Tasks::Import.bulk_import_model(index, klass, options)

  end

  task :all do
    dir         = ENV['DIR'].to_s != '' ? ENV['DIR'] : Rails.root.join("app/models")
    batch_size  = ENV['BATCH_SIZE'].to_i || 1000
    suffix      = ENV['SUFFIX'] || ''

    puts "[IMPORT] Loading models from: #{dir}"
    Dir.glob(File.join("#{dir}/**/*.rb")).each do |path|
      require path

      model_filename = path[/#{Regexp.escape(dir.to_s)}\/([^\.]+).rb/, 1]
      next if model_filename.match(/^concerns\//i) # Skip concerns/ folder
      klass          = model_filename.camelize.constantize

      # Skip if the class doesn't have Tire integration
      next unless klass.respond_to?(:tire)

      total  = klass.count rescue nil
      if suffix
        index = Tire::Index.new("#{klass.tire.index.name.to_s}_#{suffix}")
      else
        index = klass.tire.index
      end

      includes = klass.respond_to?(:includes_for_index) ? klass.includes_for_index : []
      options  = {batch_size: batch_size, include: includes}

      Tire::Tasks::Import.progress_bar(klass, total) if total
      Tire::Tasks::Import.delete_index(index) if ENV['FORCE']
      Tire::Tasks::Import.create_index(index, klass)
      puts "[IMPORT] #{klass.to_s} Records:   #{total}"
      Tire::Tasks::Import.bulk_import_model(index, klass, options)
      puts
    end

    puts '[IMPORT] Done.'
  end
end