mongoimport.rb

require "open3"

# Helper for calling mongoimport utility
module MongoImport
  include KoanHealth::Logging

  MONGO_IMPORT_BIN = 'mongoimport'
  DEFAULT_OPTIONS = {
      format: :json,
      header: false,
      drop: false,
      upsert: false,
      upsert_fields: [],
      stop_on_error: true,
      ignore_blanks: false,
      auth_db: 'admin'
  }

  # Imports a collection
  # Options
  #   format:         :json (default), :csv, :tsv
  #   fields:         comma separated list of field names (use if no header)
  #   header:         default false, use first line as headers
  #   upsert:         default false, insert or update objects that already exist
  #   upsert_fields:  default [], comma-separated fields for the query part of the upsert
  #   drop:           default false, drop collection first
  #   stop_on_error:  default true, stop importing at first error rather than continuing
  #   ignore_blanks:  default false, empty fields in csv and tsv will be ignored
  #
  def self.import_collection(mongo_session, collection, file, options = { })
    cmd = build_mongoimport_cmd(mongo_session, collection, file, DEFAULT_OPTIONS.merge(options))

    logger.info hide_user_and_password("EXEC: #{cmd}")
    stdout, stderr, exit_status = Open3.capture3(cmd)
    logger.debug "  OUT: #{stdout}"
    logger.debug "  ERR: #{stderr}"
    fail "mongoimport failed: #{stderr}" unless exit_status.success?
  end

  private

  def self.build_mongoimport_cmd(mongo_session, collection, file, options)
    # TODO: should use current session, not just the default
    file_path = file.respond_to?(:path) ? file.path : file.to_s
    cmd = ''
    cmd << MONGO_IMPORT_BIN
    cmd << " -h #{MongoUtil.get_host(mongo_session)}"
    cmd << " -d #{mongo_session.options[:database]}"
    cmd << " -u #{Mongoid.sessions[:default][:username]}" if Mongoid.sessions[:default][:username]
    cmd << " -p #{Mongoid.sessions[:default][:password]}" if Mongoid.sessions[:default][:password]
    cmd << " -c #{collection.to_s}"
    cmd << " --authenticationDatabase \"#{options[:auth_db]}\""
    cmd << " --file \"#{file_path}\""
    cmd << " --type #{options[:format]}"
    cmd << ' --headerline' if options[:header]
    cmd << ' --upsert' if options[:upsert]
    cmd << " --upsertFields #{options[:upsert_fields].join(',')}" unless options[:upsert_fields].blank?
    cmd << ' --drop' if options[:drop]
    cmd << ' --stopOnError' if options[:stop_on_error]
    cmd << ' --ignoreBlanks' if options[:ignore_blanks]
    cmd << " -f #{options[:fields].join(',')}" if options[:fields]
    cmd
  end

  def self.hide_user_and_password(s)
    return s if s.blank?
    s.gsub(/(.+\s+-u\s+).+?(\s+-.+)/i, '\1********\2').gsub(/(.+\s+-p\s+).+?(\s+-.+)/i, '\1********\2')
  end

end