django-orchestra-test/orchestra/management/commands/setuppostgres.py

import os
import textwrap

from django.conf import settings
from django.core.management.base import BaseCommand, CommandError

from orchestra.utils.paths import get_project_dir
from orchestra.utils.python import random_ascii
from orchestra.utils.sys import run, check_root


class Command(BaseCommand):
    help = 'Setup PostgreSQL database.'

    def __init__(self, *args, **kwargs):
        super(Command, self).__init__(*args, **kwargs)
        # Get defaults from settings, if exists
        try:
            self.defaults = settings.DATABASES['default']
        except (AttributeError, KeyError):
            defaults = {}
        else:
            if self.defaults['ENGINE'] != 'django.db.backends.postgresql_psycopg2':
                self.defaults = {}
        
    def add_arguments(self, parser):
        parser.add_argument(
            '--db_name', 
            dest='db_name',
            default=self.defaults.get('DB_NAME', 'orchestra'),
            help='Specifies the database to create.',
            type=str
        )
        parser.add_argument(
            '--db_user', 
            dest='db_user',
            default=self.defaults.get('DB_USER', 'orchestra'),
            help='Specifies the database to create.',
            type=str
        )
        parser.add_argument(
            '--db_password', 
            dest='db_password',
            default=self.defaults.get('PASSWORD', ''),
            help='Specifies the database password, random if not specified.',
            type=str
        )
        parser.add_argument(
            '--db_host', 
            dest='db_host',
            default=self.defaults.get('HOST', 'localhost'),
            help='Specifies the database to create.',
            type=str
        )
        parser.add_argument(
            '--db_port', 
            dest='db_port',
            default=self.defaults.get('PORT', '5432'),
            help='Specifies the database to create.',
            type=str
        )
        parser.add_argument(
            '--noinput', 
            action='store_false',
            dest='interactive',
            default=True,
            help='''Tells Django to NOT prompt the user for input of any kind. 
                 You must use --username with --noinput, and must contain the 
                 cleeryd process owner, which is the user how will perform tincd updates'''
        )
    
    def run_postgres(self, cmd, *args, **kwargs):
        return run('su postgres -c "psql -c \\"%s\\""' % cmd, *args, **kwargs)
    
    @check_root
    def handle(self, *args, **options):
        interactive = options.get('interactive')
        db_password = options.get('db_password')
        context = {
            'db_name': options.get('db_name'),
            'db_user': options.get('db_user'),
            'db_password': db_password,
            'db_host': options.get('db_host'),
            'db_port': options.get('db_port'),
            'default_db_password': db_password or random_ascii(10),
        }
        
        create_user = "CREATE USER %(db_user)s PASSWORD '%(default_db_password)s';"
        alter_user = "ALTER USER %(db_user)s WITH PASSWORD '%(db_password)s';"
        create_database = "CREATE DATABASE %(db_name)s OWNER %(db_user)s;"
        
        # Create or update user
        if self.run_postgres(create_user % context, valid_codes=(0,1)).exit_code == 1:
            if interactive and not db_password:
                msg = ("Postgres user '%(db_user)s' already exists, "
                       "please provide a password [%(default_db_password)s]: " % context)
                context['db_password'] = input(msg) or context['default_db_password']
                self.run_postgres(alter_user % context)
                msg = "Updated Postgres user '%(db_user)s' password: '%(db_password)s'"
                self.stdout.write(msg % context)
            elif db_password:
                self.run_postgres(alter_user % context)
                msg = "Updated Postgres user '%(db_user)s' password: '%(db_password)s'"
                self.stdout.write(msg % context)
            else:
                raise CommandError("Postgres user '%(db_user)s' already exists and "
                                   "--db_password has not been provided." % context)
        else:
            context['db_password'] = context['default_db_password']
            msg = "Created new Postgres user '%(db_user)s' with password '%(db_password)s'"
            self.stdout.write(msg % context)
        self.run_postgres(create_database % context, valid_codes=(0,1))
        
        context.update({
            'settings': os.path.join(get_project_dir(), 'settings.py')
        })
        
        if run("grep '^DATABASES\s*=\s*{' %(settings)s" % context, valid_codes=(0,1)).exit_code == 0:
            # Update existing settings_file
            run(textwrap.dedent("""sed -i \\
                -e "s/'ENGINE':[^#]*/'ENGINE': 'django.db.backends.postgresql_psycopg2',  /" \\
                -e "s/'NAME':[^#]*/'NAME': '%(db_name)s',  /" \\
                -e "s/'USER':[^#]*/'USER': '%(db_user)s',  /" \\
                -e "s/'PASSWORD':[^#]*/'PASSWORD': '%(db_password)s',  /" \\
                -e "s/'HOST':[^#]*/'HOST': '%(db_host)s',  /" \\
                -e "s/'PORT':[^#]*/'PORT': '%(db_port)s',  /" %(settings)s\
                """) % context
            )
        else:
            db_config = textwrap.dedent("""\
                DATABASES = {
                    'default': {
                        'ENGINE': 'django.db.backends.postgresql_psycopg2',
                        'NAME': '%(db_name)s',
                        'USER': '%(db_user)s',
                        'PASSWORD': '%(db_password)s',
                        'HOST': '%(db_host)s',
                        'PORT': '%(db_port)s',
                        'ATOMIC_REQUESTS': True,
                    }
                }""") % context
            context.update({
                'db_config': db_config
            })
            run('echo "%(db_config)s" >> %(settings)s' % context)