db.py

"""
DB migrations
"""

import argparse
import os
import shellish
import subprocess
import psycopg2
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT


class DB(shellish.Command):
    """ Perform DB management. """

    name = 'db'

    def setup_args(self, parser):
        self.add_subcommand(Create)
        self.add_subcommand(Remove)
        self.add_subcommand(Migrate)
        self.add_subcommand(Rollback)
        self.add_subcommand(Status)


class DBCommand(shellish.Command):

    def setup_args(self, parser):
        self.add_argument('--host', env='DATABASE_HOST',
                          help="Database server host or IP.")
        self.add_argument('--port', type=int, env='DATABASE_PORT',
                          help="Database server port.")
        self.add_argument('--user', env='DATABASE_USER',
                          help="Database username.")
        self.add_argument('--password', env='DATABASE_PASSWORD',
                          help="Database password.")
        self.add_argument('--dbname', env='DATABASE_NAME',
                          help="Database name to create.")

    def sqlexec(self, args, statement):
        with psycopg2.connect(host=args.host, port=args.port, user=args.user,
                              password=args.password, dbname='postgres') as db:
            db.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
            with db.cursor() as c:
                c.execute(statement)


class Create(DBCommand):
    """ Create a new database for this stack. """

    name = 'create'

    def run(self, args):
        raise ValueError("DEPRECATED COMMAND")
        shellish.vtmlprint("<b>Creating database:</b> <blue>%s</blue>" %
                           args.dbname)
        try:
            self.sqlexec(args, 'CREATE DATABASE %s;' % args.dbname)
        except psycopg2.ProgrammingError as e:
            if 'already exists' in str(e):
                shellish.vtmlprint("<b>WARNING:</b> <yellow>Database already exists")
            else:
                raise e


class Remove(DBCommand):
    """ Create a new database for this stack. """

    name = 'rm'

    def setup_args(self, parser):
        super().setup_args(parser)
        self.add_argument('--force', '-f', action='store_true',
                          help="Do not prompt for removal")

    def run(self, args):
        raise ValueError("DEPRECATED COMMAND")
        if not args.force:
            confirm = input("Type name of db (%s) to confirm removal: " %
                            args.dbname)
            if confirm != args.dbname:
                raise SystemExit("Aborted database removal")
        shellish.vtmlprint("<b>WARNING:</b> <red>Removing database <b>%s</b>"
                           % args.dbname)
        self.sqlexec(args, 'DROP DATABASE %s;' % args.dbname)


class PhinxCommand(shellish.Command):

    phinx_command = None

    def setup_args(self, parser):
        self.add_argument('args', nargs=argparse.REMAINDER,
                          help='Args for phinx command.')

    def run(self, args):
        cmd = [
            '%s/bin/phinx' % os.environ['COMPOSER_VENDOR_DIR'],
            self.phinx_command
        ]
        cmd.extend(args.args)
        raise SystemExit(subprocess.call(cmd, cwd='/db'))


class Migrate(PhinxCommand):
    """ Perform DB migrations with Phinx. """

    name = 'migrate'
    phinx_command = 'migrate'

class Rollback(PhinxCommand):
    """ Perform DB migrations with Phinx. """

    name = 'rollback'
    phinx_command = 'rollback'


class Status(PhinxCommand):
    """ Show DB schema status from Phinx. """

    name = 'status'
    phinx_command = 'status'