pg_backup.py

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

from optparse import OptionParser
from datetime import date, datetime
from pwd import getpwnam
import unidecode
from libs.script_utils import get_email_logger

import getpass
import subprocess
import stat
import settings

admins = [item[1] for item in settings.ADMINS]
email_settings = {
    'user': settings.EMAIL_HOST_USER,
    'host': settings.EMAIL_HOST,
    'password': settings.EMAIL_HOST_PASSWORD,
}
logger = get_email_logger(admins, u"Backup error", email_settings)

parser = OptionParser()
parser.add_option("-L", "--lockfile", dest="lockfile", help="lock file",
                    default="/tmp/mf_backup_script.lock")
parser.add_option("-d", "--database", dest="database", help="database name",
                    default="finance")
parser.add_option("-t", "--db", dest="db_type", help="database type",
                    default="postgresql")
parser.add_option("-f", "--folder", dest="folder", help="backup folder",
                    default="/home/%s/backup" % getpass.getuser())
parser.add_option("-c", "--compress",
                  action="store_true", dest="compress", default=False,
                  help="compress and delete uncompressed dump")
(options, args) = parser.parse_args()

import os
if not os.path.exists(options.folder):
    logger.error("Output folder does not exist")
    exit(1)

if os.path.exists(options.lockfile):
    logger.error("""Lock file exists (%s), exiting.
                     Please remove lock file, if you want to run script""" % options.lockfile)
    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'
    user = 'undefined'
    __metaclass__ = BRegister
    command = ['dir', '/home/']  # sample arguments of backup command
    stdout = False

    def init(self):
        os.chmod(options.folder, stat.S_IWOTH | stat.S_IXOTH | stat.S_IROTH \
                         | stat.S_IWUSR | stat.S_IXUSR | stat.S_IRUSR)
        os.setuid(getpwnam(self.user).pw_uid)
        open(options.lockfile, 'w')

    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 = output_filename = None
        if self.stdout:
            output_folder = self.get_outputfolder()
            output_filename = os.path.join(output_folder, '%s_%s_%02d-%02d.dump'\
                             % (options.database, date.today(), ct.hour, ct.minute))
            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 = """
                pg_dump calling error:
                -----------------------------
                stderr: %s
            """ % unidecode.unidecode(stderr.decode('utf-8'))
            raise Exception(msg)

        if options.compress and output_filename:
            args = ['gzip', output_filename]
            p = subprocess.Popen(args,
                                 stdin=subprocess.PIPE,
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE, close_fds=True)
            return_code = p.wait()
            stderr = p.stderr.read()

            if return_code > 0:
                msg = """
                    compressing error:
                    -----------------------------
                    stderr: %s
                """ % unidecode.unidecode(stderr.decode('utf-8'))
                raise Exception(msg)

    def run(self):
        exit_code = 0

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

        os.remove(options.lockfile)
        exit(exit_code)

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


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


class MongoBackuper(Backuper):
    db_type = "mongodb"
    user = 'mongodb'

    @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()