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

146 lines
5.9 KiB
Python
Raw Normal View History

2014-05-08 16:59:35 +00:00
import os
2015-10-01 18:57:31 +00:00
import textwrap
2014-05-08 16:59:35 +00:00
2015-10-02 09:31:32 +00:00
from django.conf import settings
2015-10-02 13:01:01 +00:00
from django.core.management.base import BaseCommand, CommandError
2014-05-08 16:59:35 +00:00
from orchestra.utils.paths import get_project_dir
2015-10-02 09:31:32 +00:00
from orchestra.utils.python import random_ascii
2015-04-05 10:46:24 +00:00
from orchestra.utils.sys import run, check_root
2014-05-08 16:59:35 +00:00
class Command(BaseCommand):
2021-01-30 12:55:27 +00:00
help = 'Setup PostgreSQL database.'
2014-05-08 16:59:35 +00:00
def __init__(self, *args, **kwargs):
super(Command, self).__init__(*args, **kwargs)
2015-10-02 09:31:32 +00:00
# Get defaults from settings, if exists
try:
2021-01-30 12:55:27 +00:00
self.defaults = settings.DATABASES['default']
2015-10-02 09:31:32 +00:00
except (AttributeError, KeyError):
defaults = {}
else:
2021-01-30 12:55:27 +00:00
if self.defaults['ENGINE'] != 'django.db.backends.postgresql_psycopg2':
self.defaults = {}
2015-10-02 09:31:32 +00:00
2021-01-30 12:55:27 +00:00
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',
2021-01-30 13:17:18 +00:00
default=self.defaults.get('HOST', 'localhost'),
2021-01-30 12:55:27 +00:00
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,
2021-01-30 13:17:18 +00:00
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'''
2021-01-30 12:55:27 +00:00
)
2014-05-08 16:59:35 +00:00
2015-10-02 09:31:32 +00:00
def run_postgres(self, cmd, *args, **kwargs):
2015-10-02 11:36:24 +00:00
return run('su postgres -c "psql -c \\"%s\\""' % cmd, *args, **kwargs)
2015-10-02 09:31:32 +00:00
2014-05-08 16:59:35 +00:00
@check_root
def handle(self, *args, **options):
2015-10-02 09:31:32 +00:00
interactive = options.get('interactive')
db_password = options.get('db_password')
2014-05-08 16:59:35 +00:00
context = {
2021-01-30 13:17:18 +00:00
'db_name': options.get('db_name'),
'db_user': options.get('db_user'),
2015-10-02 09:31:32 +00:00
'db_password': db_password,
2021-01-30 13:17:18 +00:00
'db_host': options.get('db_host'),
'db_port': options.get('db_port'),
2015-10-02 09:31:32 +00:00
'default_db_password': db_password or random_ascii(10),
2015-10-01 18:57:31 +00:00
}
2014-05-08 16:59:35 +00:00
2015-10-02 09:46:33 +00:00
create_user = "CREATE USER %(db_user)s PASSWORD '%(default_db_password)s';"
2015-10-02 09:31:32 +00:00
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:
2015-10-02 09:46:33 +00:00
if interactive and not db_password:
2015-10-02 09:31:32 +00:00
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']
2015-10-02 09:37:25 +00:00
self.run_postgres(alter_user % context)
msg = "Updated Postgres user '%(db_user)s' password: '%(db_password)s'"
2015-10-02 09:46:33 +00:00
self.stdout.write(msg % context)
elif db_password:
2015-10-02 09:37:25 +00:00
self.run_postgres(alter_user % context)
msg = "Updated Postgres user '%(db_user)s' password: '%(db_password)s'"
2015-10-02 09:46:33 +00:00
self.stdout.write(msg % context)
2015-10-02 09:31:32 +00:00
else:
raise CommandError("Postgres user '%(db_user)s' already exists and "
2021-01-28 18:57:42 +00:00
"--db_password has not been provided." % context)
2015-10-02 09:46:33 +00:00
else:
2015-10-02 11:20:52 +00:00
context['db_password'] = context['default_db_password']
msg = "Created new Postgres user '%(db_user)s' with password '%(db_password)s'"
2015-10-02 09:46:33 +00:00
self.stdout.write(msg % context)
self.run_postgres(create_database % context, valid_codes=(0,1))
2015-10-02 09:31:32 +00:00
2015-10-01 18:57:31 +00:00
context.update({
'settings': os.path.join(get_project_dir(), 'settings.py')
})
2014-05-08 16:59:35 +00:00
2015-10-01 18:57:31 +00:00
if run("grep '^DATABASES\s*=\s*{' %(settings)s" % context, valid_codes=(0,1)).exit_code == 0:
2014-05-08 16:59:35 +00:00
# Update existing settings_file
2015-10-02 09:31:32 +00:00
run(textwrap.dedent("""sed -i \\
2015-10-05 13:31:08 +00:00
-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\
2015-10-01 18:57:31 +00:00
""") % context
)
2014-05-08 16:59:35 +00:00
else:
2015-10-01 18:57:31 +00:00
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
2015-10-02 09:31:32 +00:00
context.update({
'db_config': db_config
})
2014-05-08 16:59:35 +00:00
run('echo "%(db_config)s" >> %(settings)s' % context)