django-orchestra-test/orchestra/management/commands/setuppostgres.py
2021-01-30 14:59:49 +01:00

146 lines
5.9 KiB
Python

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)