namespace :db do
  namespace :data do
    desc "Dump data into sql script file: filename=[target filename]"
    task :dump => 'environment' do
      environment = (ENV.include?("RAILS_ENV")) ? (ENV["RAILS_ENV"]) : 'development'
      ENV["RAILS_ENV"] = RAILS_ENV = environment

      database = get_database(environment)
      user = database
      password = ENV['PASSWORD']
      filename = ENV['FILENAME'] || environment
      timestamp = ENV['TIMESTAMP'] ? bool(ENV['TIMESTAMP']) : true
      file_suffix = DateTime.now.strftime('%Y-%m-%d-%H-%M-%S')

      full_filename = filename + (timestamp ? "_#{file_suffix}" : '') + '.sql'

      puts "Connecting to #{environment} environment..."
      sh "mysqldump -u#{user} -p#{password} #{database} --skip-triggers --compact " +
        "--ignore-table=#{database}.schema_migrations " +
        "--no-create-info > #{RAILS_ROOT}/db/data/#{full_filename}"

      puts "Successfully created #{full_filename} from #{environment} environment."
    end

    desc "Load data from sql script file: filename=[data file]"
    task :load => 'environment' do
      environment = (ENV.include?("RAILS_ENV")) ? (ENV["RAILS_ENV"]) : 'development'
      ENV["RAILS_ENV"] = RAILS_ENV = environment

      database = get_database(environment)
      user = database
      password = ENV['PASSWORD']
      filename = ENV['FILENAME']
      raise "Please specify a source file (FILENAME=[source.sql])" if filename.blank?

      puts "Connecting to #{environment}..."
      ActiveRecord::Base.establish_connection(RAILS_ENV.to_sym)

      puts "Truncating tables..."
      ActiveRecord::Base.connection.execute('show tables').each do |table|
        unless table.to_s == 'schema_migrations'
          puts "   Truncating #{table}"
          ActiveRecord::Base.connection.execute("truncate table #{table.to_s}")
        end
      end

      puts "Importing data from #{filename}..."
      sh "mysql -u#{user} -p#{password} #{database} < #{RAILS_ROOT}/#{filename}"
      puts "Completed loading #{filename} into #{environment} environment."
    end
  end
end

private
def get_database(environment)
  case environment
    when 'test'
      raise 'add your TEST db name'
    when 'staging'
      raise 'add your STAGING db name'
    when 'production'
      raise 'add your PRODUCTION db name'
    else
      raise 'add your DEVELOPMENT db name'
  end
end

def bool(s)
  s.match(/(true|t|yes|y|1)$/i) != nil
end