walter/lib/backup-push.js

Backup utilities for postgresql

var env = process.env;
var fs = require('fs');
var pg = require('pg.js');
var path = require('path');
var knox = require('knox');
var moment = require('moment');
var es = require('event-stream');
var buffers = require('buffers');
var queries = require('./queries');
var optimist = require('optimist');
var format = require('util').format;
var mkdirp = require('mkdirp').sync;
var spawn = require('child_process').spawn;
var winston = require('winston');
var sbuffers = require('stream-buffers');

var argv = optimist
  .usage('Push a base backup to s3.\nUsage: walter backup-push [options]')
  .describe('pg-dir', 'Postgres data directory (env: WALTER_PG_DIR)')
  .describe('pg-conn', 'Postgres connection string (env: WALTER_PG_CONN)')
  .describe('working-dir', 'Working directory for walter (env: WALTER_DIR)')
  .describe('prefix', 'Prefix used to store files on s3 (env: WALTER_PREFIX)')
  .describe('gpg-id', 'GPG ID to encrypt your backups (env: WALTER_GPG_ID)')
  .describe('aws-key', 'AWS Key (env: WALTER_AWS_KEY)')
  .describe('aws-secret', 'AWS Secret (env: WALTER_AWS_SECRET)')
  .describe('aws-region', 'AWS region for your bucket (env: WALTER_AWS_REGION)')
  .describe('s3-bucket', 'S3 Bucket Name (env: WALTER_S3_BUCKET)')
  .describe('s3-part-size', 'S3 part size for multi-part uploads (env: WALTER_S3_PART_SIZE')
  .describe('log-file-path', 'Winston file transport path (env: WALTER_LOG_FILE_PATH)')
  .describe('log-file-size', 'Winston file transport size in MB (env: WALTER_LOG_FILE_SIZE)')
  .describe('log-file-limit', 'Winston file transport rotate limit (env: WALTER_LOG_FILE_LIMIT)')
  .describe('labs-api-key', 'Walter Labs API Key (env: WALTER_LABS_API_KEY)')
  .describe('log-http-host', 'Winston HTTP transport host (env: WALTER_LOG_HTTP_HOST)')
  .describe('log-http-port', 'Winston HTTP transport port (env: WALTER_LOG_HTTP_PORT)')
  .describe('log-http-path', 'Winston HTTP transport path (env: WALTER_LOG_HTTP_PATH)')
  .describe('log-http-ssl', 'Winston HTTP transport ssl (env: WALTER_LOG_HTTP_SSL)')
  .describe('tar-bin', 'Path to tar binary')
  .describe('gpg-bin', 'Path to gpg binary')
  .describe('lzop-bin', 'Path to lzop binary')
  .describe('upload-chunk-size', 'Size of chunk for streaming upload')
  .describe('upload-chunk-frequency', 'Frequency between streaming upload chunks')
  .default('pg-dir', env.WALTER_PG_DIR)
  .default('pg-conn', env.WALTER_PG_CONN || 'postgres://postgres@localhost/postgres')
  .default('working-dir', env.WALTER_DIR || path.join(getUserHome(), '.walter'))
  .default('prefix', env.WALTER_PREFIX)
  .default('gpg-id', env.WALTER_GPG_ID)
  .default('aws-key', env.WALTER_AWS_KEY)
  .default('aws-secret', env.WALTER_AWS_SECRET)
  .default('s3-bucket', env.WALTER_S3_BUCKET)
  .default('s3-part-size', env.WALTER_S3_PART_SIZE || (5*1024*1024))
  .default('log-file-path', env.WALTER_LOG_FILE_PATH)
  .default('log-file-size', env.WALTER_LOG_FILE_SIZE)
  .default('log-file-limit', env.WALTER_LOG_FILE_LIMIT)
  .default('labs-api-key', env.WALTER_LABS_API_KEY)
  .default('log-http-host', env.WALTER_LOG_HTTP_HOST)
  .default('log-http-port', env.WALTER_LOG_HTTP_PORT)
  .default('log-http-path', env.WALTER_LOG_HTTP_PATH)
  .default('log-http-ssl', env.WALTER_LOG_HTTP_SSL)
  .default('upload-chunk-size', 1048576)
  .default('upload-chunk-frequency', 1)
  .default('tar-bin', 'tar')
  .default('gpg-bin', 'gpg')
  .default('lzop-bin', 'lzop')
  .demand('pg-dir')
  .demand('prefix')
  .demand('aws-key')
  .demand('aws-secret')
  .demand('s3-bucket')
  .wrap(120)
  .argv;

winston.remove(winston.transports.Console);

if(argv['labs-api-key']) {
  argv['log-http-host'] = 'api.walter-labs.com';
  argv['log-http-port'] = 443;
  argv['log-http-path'] = '/servers/' + argv['labs-api-key'];
  argv['log-http-ssl'] = true;
}

if(argv['log-http-host']) {
  var logHttpOpts = {
    host: argv['log-http-host']
  };

  if(argv['log-http-port'])
    logHttpOpts.port = argv['log-http-port'];

  if(argv['log-http-path'])
    logHttpOpts.path = argv['log-http-path'];

  if(argv['log-http-ssl'] === 'true' || argv['log-http-ssl'] === true)
    logHttpOpts.ssl = true;

  winston.add(winston.transports.Http, logHttpOpts);
}

if(argv['log-file-path']) {
  var logFileOpts = {
    filename: argv['log-file-path']
  };

  if(argv['log-file-size'])
    logFileOpts.maxsize = (argv['log-file-size']*1024*1024);

  if(argv['log-file-limit'])
    logFileOpts.maxFiles = argv['log-file-limit'];

  winston.add(winston.transports.File, logFileOpts);
}

var backup_running = false;
var mpu_running = false;
var db = new pg.Client(argv['pg-conn']);
var label = 'backup_started_' + moment().utc().toJSON();
var s3 = createS3Client();
var started = new Date();

winston.info('base-backup-push: %s: started.', argv['prefix'], {
  type: 'base',
  action: 'started',
  prefix: argv['prefix'],
  started: started
});

var rl = require('readline').createInterface({
  input: process.stdin,
  output: process.stdout
});

rl.setPrompt('', 0);
rl.clearLine = function() {
  this.write(null, {ctrl: true, name: 'u'});
};

module.exports = function() {
  db.connect(function(err) {
    if(err) throw err;
    db.query(queries.start, [label], onStarted);
  });
};

function onStarted(err, data) {
  if(err) return handleError(err);

  var start_info = data.rows[0];
  var base_path = argv['prefix'] + '/base/';
  var base_file = start_info.offset + '_' + start_info.file_name;
  var manifest_file = base_file + '_manifest.json';
  var working_path = path.resolve(argv['working-dir'], base_path);
  var working_file = path.resolve(working_path, base_file);
  var remote_path = base_path + base_file;

  backup_running = true;

  console.log('Initiated PostgreSQL base backup:', base_file);
  console.log('Initiating S3 MPU upload:', remote_path);

  initiateMPU(remote_path, function(err, mpu_id) {
    if(err) return handleError(err);

    console.log('Initiated S3 MPU upload:', mpu_id);
    console.log('Initiating streaming part uploads.');

    mpu_running = true;
    mkdirp(working_path);

    var tar = spawnTar();
    var lzop = spawnLzop();
    var gpg = spawnGpg();
    var ps = es.pause();
    var chunker = createChunker();
    var uploader = createUploader(remote_path, mpu_id, ps, 5);
    var writer = createArrayWriter(onUploadComplete);

    var pipeline = gpg
      ? [tar, lzop, gpg, ps, chunker, uploader, writer]
      : [tar, lzop, ps, chunker, uploader, writer];

    var pipe = es.pipeline.apply(null, pipeline);
    pipe.on('error', handleError);

    function onUploadComplete(err, parts) {
      if(err) return handleError(err);

      rl.resume();
      rl.close();

      console.log('Succesfully uploaded', parts.length, 'parts');

      completeMPU(remote_path, mpu_id, parts, function(err) {
        if(err) return handleError(err);
        console.log('Completed S3 MPU upload:', remote_path);
        mpu_running = false;
        db.query(queries.stop, onStopped);
      });
    }
  });

  function onStopped(err, result) {

    if(err) return handleError(err);

    console.log('Completed PostgreSQL base backup:', result.rows[0]);

    backup_running = false;
    var manifest = createManifest(result.rows[0]);
    var manifestSize = Buffer.byteLength(manifest);

    console.log('Uploading manifest to S3');

    var putManifest = s3.put(base_path + manifest_file, {
      'Content-Length': manifestSize,
      'Content-Type': 'application/json'
    });

    putManifest.on('response', manifestResponse);
    putManifest.end(manifest, 'utf8');
  }

  function manifestResponse(response) {
    if((response.statusCode < 200 || response.statusCode > 299)) {
      var err = new Error('Failed manifest upload to S3');
      return handleError(err);
    }

    console.log('Successfully uploaded manifest');

    s3.headFile(remote_path, function(err, res) {

      if(err || (res.statusCode < 200 || res.statusCode > 299)) {
        if(!err) err = new Error('Failed to retrieve head for base backup after upload');
        return handleError(err);
      }

      var filesize = parseInt(res.headers['content-length'], 10);

      winston.info('base-backup-push: %s: completed.', argv['prefix'], {
        type: 'base',
        action: 'completed',
        prefix: argv['prefix'],
        base_offset: start_info.offset,
        base_file_name: start_info.file_name,
        manifest: manifest_file,
        filename: base_file,
        started: started,
        size: filesize,
        completed: new Date()
      }, function() {
        db.end();
      });

    });
  }
}

function initiateMPU(path, callback) {
  var req = s3.request('POST', path + '?uploads');

  req.on('response', function(res, body) {
    var xml = '';

    res.on('data', function(data) {
      xml += data;
    });

    res.on('end', function() {
      var parsed = xml.match(/\<UploadId\>(.*)\<\/UploadId\>/);

      if((res.statusCode < 200 || res.statusCode > 299) || !parsed || parsed.length !== 2) {
        var err = new Error('Unable to initiate multi-part upload.');
        return callback(err);
      }

      callback(null, parsed[1]);
    });

    res.on('error', callback);
  });

  req.end();
}

function completeMPU(path, mpu_id, parts, callback) {

  parts.sort(sortByPartNumber);

  var xml = '';
  var tpl = '<Part><PartNumber>%d</PartNumber><ETag>%s</ETag></Part>';

  for(var i=0; i<parts.length; i++)
    xml += format(tpl, parts[i].number, parts[i].etag);

  var req = s3.request('POST', path + '?uploadId=' + mpu_id);

  req.on('response', function(res, body) {
    var response = '';

    res.on('data', function(data) {
        response += data;
    });

    res.on('end', function() {
      var err = null;

      if((res.statusCode < 200 || res.statusCode > 299))
        err = new Error(response);

      callback(err);
    });

    res.on('error', callback);
  });

  req.write('<CompleteMultipartUpload>' + xml + '</CompleteMultipartUpload>');
  req.end();
}

function createChunker() {
  var bufs = buffers();

  function onData(data) {
    bufs.push(data);
    if(bufs.length >= argv['s3-part-size']) {
      this.queue(bufs.toBuffer());
      bufs = buffers();
    }
  }

  function onEnd() {
    if(bufs.length > 0) {
      this.queue(bufs.toBuffer());
      bufs = null;
    }
    this.queue(null);
  }

  return es.through(onData, onEnd);
}

function createUploader(remote_path, mpu_id, ps, concurrency) {
  var active = 0;
  var part = 0;
  var completed = 0;
  var completedBytes = 0;
  var status = 'Uploading parts, %d active uploads, %d completed (%s)';
  var started = false;

  function updateStatus() {
     var mbs = (completedBytes/1024/1024).toFixed(2) + 'MB';
     if(started) rl.clearLine();
     if(!started) started = true;
     rl.write(format(status, active, completed, mbs));
  }

  return es.map(function(data, callback) {
    active++;

    var attempts = 1;
    var number = ++part;
    var url = remote_path + '?partNumber=' + number + '&uploadId=' + mpu_id;
    var headers = {'Content-Length': data.length};
    if(active === concurrency) ps.pause();


    function handleResponse(err, res) {
      if(err || !res || ((res.statusCode < 200 || res.statusCode > 299))) {
        if(attempts >= 3) return callback(new Error('Failed uploading part.'));
        return retryPart();
      }

      active--;
      completed++;
      completedBytes += data.length;
      updateStatus();

      if(active < concurrency && ps.paused)
        ps.resume();

      var etag = res.headers.etag.slice(1, res.headers.etag.length-1);
      return callback(null, {etag: etag, number: number});
    }

    function retryPart() {
      attempts++;
      uploadPart(url, headers, data, handleResponse);
    }

    updateStatus();
    uploadPart(url, headers, data, handleResponse);
  });
}

function uploadPart(url, headers, data, callback) {

  var req = s3.request('PUT', url, headers);

  var sbuffer = new sbuffers.ReadableStreamBuffer({
    frequency: argv['upload-chunk-frequency'],
    chunkSize: argv['upload-chunk-size']
  });

  sbuffer.put(data);

  req.on('error', callback);
  req.on('response', function(res) {
    return callback(null, res);
  });

  sbuffer.pipe(req);
}


function createArrayWriter(callback) {
  return es.writeArray(callback);
}

function putBase(working_file, base_path, base_file, manifest_file) {

  var stat = fs.statSync(working_file);
  var rs = fs.createReadStream(working_file);

  var put = s3.put(base_path + base_file, {
    'Content-Length': stat.size
  });

  put.on('error', handleError);
  put.on('response', baseResponse);
  rs.pipe(put);

  function baseResponse(response) {
    if((response.statusCode < 200 || response.statusCode > 299)) {
      var err = new Error('Failed upload to S3');
      return handleError(err);
    }

    db.query(queries.stop, onStopped);
  }

  function onStopped(err, result) {

    if(err) return handleError(err);

    backup_running = false;
    var manifest = createManifest(result.rows[0]);
    var manifestSize = Buffer.byteLength(manifest);

    var putManifest = s3.put(base_path + manifest_file, {
      'Content-Length': manifestSize,
      'Content-Type': 'application/json'
    });

    putManifest.on('response', manifestResponse);
    putManifest.end(manifest, 'utf8');
  }

  function manifestResponse(response) {
    if((response.statusCode < 200 || response.statusCode > 299)) {
      var err = new Error('Failed manifest upload to S3');
      return handleError(err);
    }

    fs.unlinkSync(working_file);
    db.end();
  }
}

function createManifest(stop_info) {
  return JSON.stringify({
      'wal_segment_name': stop_info.file_name,
      'wal_segment_offset': stop_info.file_offset
  });
}

function createS3Client() {
  var options = {
    key: argv['aws-key'],
    secret: argv['aws-secret'],
    bucket: argv['s3-bucket']
  };

  if(argv['aws-region'])
    options.region = argv['aws-region'];

  return new knox.createClient(options);
}

function spawnTar() {
  var args = [
    '--exclude', '**/pg_xlog/*',
    '--exclude', '**/pg_log/*',
    '--exclude', '**/postmaster.pid',
    '--exclude', '**/postmaster.opts',
    '--exclude', '**/postgresql.conf',
    '-c', '-C', argv['pg-dir'], './'
  ];
  return es.child(spawn(argv['tar-bin'], args));
}

function spawnLzop() {
  return es.child(spawn(argv['lzop-bin'], ['-c']));
}

function spawnGpg() {
  if(!argv['gpg-id']) return false;
  var args = ['-e', '-z', '0', '-r', argv['gpg-id'], '-'];
  return es.child(spawn(argv['gpg-bin'], args));
}

function onConnect(err) {
  if(err) handleError(err);
}

function handleError(err) {
  winston.error('base-backup-push: %s: error occured.', argv['prefix'], {
    prefix: argv['prefix'],
    error: err.message,
    started: started,
    errored: new Date()
  }, function() {
    if(!backup_running) {
      db.end();
      throw err;
    }

    db.query(queries.stop, function() {
      db.end();
      throw err;
    });
  });
}

function getUserHome() {
  return process.env.HOME || process.env.HOMEPATH || process.env.USERPROFILE;
}

function sortByPartNumber(a, b) {
  if (a.number < b.number)
     return -1;
  if (a.number > b.number)
    return 1;
  return 0;
}