pg_backup.py

#!/usr/bin/env python
#coding: utf-8

from optparse import OptionParser
from datetime import date, datetime

import unidecode
import subprocess
import fcntl
import os
import sys
import logging
import logging.handlers

parser = OptionParser()
parser.add_option("-d", "--database", dest="database", help="database name")
parser.add_option("-t", "--type", dest="db_type", help="database type",
                    default="postgresql")
parser.add_option("-f", "--folder", dest="folder", help="backup folder")
parser.add_option("-c", "--compress",
                  action="store_true", dest="compress", default=False,
                  help="compress and delete uncompressed dump")
(options, args) = parser.parse_args()

if not options.database:
    print "Specify database name"
    exit(1)

if not options.folder or not os.path.exists(options.folder):
    print "Output folder does not exist"
    exit(1)

logging.basicConfig()
log_filename = '/tmp/pg_backup_%s.log' % options.database
log_file_handler = logging.handlers.RotatingFileHandler(log_filename,
                                    maxBytes=1024 * 1024 * 2, backupCount=2)
logger = logging.getLogger(options.database)
logger.addHandler(log_file_handler)
logger.setLevel(logging.DEBUG)

lockfile = "/tmp/pg_backup_py_%s.lock" % options.database
fp = open(lockfile, 'w')
try:
    fcntl.lockf(fp, fcntl.LOCK_EX | fcntl.LOCK_NB)
except IOError:
    logger.error("Another instance of pg_backup process is already running, quitting.")
    sys.exit(-1)


BACKUPERS = {}


class BRegister(type):
    def __new__(self, *args, **kwargs):
        backuper = type.__new__(self, *args, **kwargs)
        BACKUPERS[backuper.db_type] = backuper

        return backuper


class Backuper(object):
    db_type = 'undefined'
    __metaclass__ = BRegister
    command = ['dir', '/home/']  # sample arguments of backup command
    extension = 'dump'

    def init(self):
        pass

    def get_outputfolder(self, make=True):
        folder = os.path.join(options.folder, options.database, '%s/' % date.today())
        if make and not os.path.exists(folder):
            os.makedirs(folder)

        return folder

    def backup(self):
        ct = datetime.now().time()
        output_file = None
        output_folder = self.get_outputfolder()
        output_filename = os.path.join(output_folder, '%s_%s_%02d-%02d.%s'
                         % (options.database, date.today(), ct.hour, ct.minute,
                            self.extension))

        if options.compress and self.compress_command:
            output_filename += '.gz'
            output_file = open(output_filename, 'w')
            dumper = subprocess.Popen(self.compress_command,
                                stdout=subprocess.PIPE)
            p = subprocess.Popen(['gzip'],
                                 stdin=dumper.stdout,
                                 stderr=subprocess.PIPE, close_fds=True,
                                 stdout=output_file or subprocess.PIPE)
        else:
            output_file = open(output_filename, 'w')
            p = subprocess.Popen(self.command,
                                 stdin=subprocess.PIPE,
                                 stdout=output_file or subprocess.PIPE,
                                 stderr=subprocess.PIPE, close_fds=True)

        return_code = p.wait()
        stderr = p.stderr.read()

        if output_file:
            output_file.close()

        if return_code > 0:
            if output_filename:
                os.remove(output_filename)

            msg = """
                dump command calling error:
                -----------------------------
                stderr: %s
            """ % unidecode.unidecode(stderr.decode('utf-8'))
            raise Exception(msg)
        else:
            logger.info("%s ok" % output_filename)

    def run(self):
        exit_code = 0

        try:
            self.init()
            self.backup()
        except Exception:
            info = {'db_type': options.db_type,
                    'folder': options.folder}
            logger.error("Error occured during backup, info = %s" % info, exc_info=1)
            exit_code = 1

        exit(exit_code)

    @classmethod
    def get_backuper(cls, db_type):
        return BACKUPERS[db_type]


class PostgresBackuper(Backuper):
    db_type = 'postgresql'
    command = ['pg_dump', '-Fc', options.database]
    compress_command = ['pg_dump', options.database]


class MongoBackuper(Backuper):
    db_type = "mongodb"
    compress_command = None
    extension = 'log'

    @property
    def command(self):
        output_folder = self.get_outputfolder()
        return ['mongodump', '-d', options.database, '-o', output_folder]


if __name__ == '__main__':
    backuper = Backuper.get_backuper(options.db_type)
    backuper().run()