class PostgresTest(unittest.TestCase):
"""Test the public API provided by the `Postgres` class"""
def setUp(self):
"""Set up a temporary database cluster for testing potentially destructive operations"""
self.pg_test = PGTest()
self.postgres = Postgres(port=self.pg_test.port,
interactive=False,
quiet=True)
self.dbuser = '******'
self.dbpass = '******'
self.dbname = 'aiida_db'
def _setup_postgres(self):
self.postgres.dbinfo = self.pg_test.dsn
self.postgres.determine_setup()
def test_determine_setup_fail(self):
self.postgres.set_port('11111')
setup_success = self.postgres.determine_setup()
self.assertFalse(setup_success)
def test_determine_setup_success(self):
self._setup_postgres()
self.assertTrue(self.postgres.pg_execute)
def test_setup_fail_callback(self):
"""Make sure `determine_setup` works despite wrong initial values in case of correct callback"""
def correct_setup(interactive, dbinfo): # pylint: disable=unused-argument
return self.pg_test.dsn
self.postgres.set_port(11111)
self.postgres.set_setup_fail_callback(correct_setup)
self.postgres.determine_setup()
self.assertTrue(self.postgres.pg_execute)
@mock.patch('aiida.control.postgres._try_connect',
new=_try_connect_always_fail)
@mock.patch('aiida.control.postgres._try_subcmd')
def test_fallback_on_subcmd(self, try_subcmd):
"""Ensure that accessing postgres via subcommand is tried if psychopg does not work."""
self._setup_postgres()
self.assertTrue(try_subcmd.call_count >= 1)
def test_create_drop_db_user(self):
"""Check creating and dropping a user works"""
self._setup_postgres()
self.postgres.create_dbuser(self.dbuser, self.dbpass)
self.assertTrue(self.postgres.dbuser_exists(self.dbuser))
self.postgres.drop_dbuser(self.dbuser)
self.assertFalse(self.postgres.dbuser_exists(self.dbuser))
def test_create_drop_db(self):
"""Check creating & destroying a database"""
self._setup_postgres()
self.postgres.create_dbuser(self.dbuser, self.dbpass)
self.postgres.create_db(self.dbuser, self.dbname)
self.assertTrue(self.postgres.db_exists(self.dbname))
self.postgres.drop_db(self.dbname)
self.assertFalse(self.postgres.db_exists(self.dbname))
def quicksetup(self, profile, email, first_name, last_name, institution,
backend, db_port, db_user, db_user_pw, db_name, repo,
set_default, non_interactive):
'''Set up a sane aiida configuration with as little interaction as possible.'''
from aiida.common.setup import create_base_dirs, AIIDA_CONFIG_FOLDER
create_base_dirs()
aiida_dir = os.path.expanduser(AIIDA_CONFIG_FOLDER)
# access postgres
postgres = Postgres(port=db_port,
interactive=bool(not non_interactive),
quiet=False)
postgres.set_setup_fail_callback(prompt_db_info)
success = postgres.determine_setup()
if not success:
sys.exit(1)
# default database name is <profile>_<login-name>
# this ensures that for profiles named test_... the database will also
# be named test_...
import getpass
osuser = getpass.getuser()
dbname = db_name or profile + '_' + osuser
# default database user name is aiida_qs_<login-name>
# default password is random
dbuser = db_user or 'aiida_qs_' + osuser
from aiida.common.setup import generate_random_secret_key
dbpass = db_user_pw or generate_random_secret_key()
# check if there is a profile that contains the db user already
# and if yes, take the db user password from there
# This is ok because a user can only see his own config files
from aiida.common.setup import (set_default_profile, get_or_create_config)
confs = get_or_create_config()
profs = confs.get('profiles', {})
for v in profs.itervalues():
if v.get('AIIDADB_USER', '') == dbuser and not db_user_pw:
dbpass = v.get('AIIDADB_PASS')
print 'using found password for {}'.format(dbuser)
break
try:
create = True
if not postgres.dbuser_exists(dbuser):
postgres.create_dbuser(dbuser, dbpass)
else:
dbname, create = _check_db_name(dbname, postgres)
if create:
postgres.create_db(dbuser, dbname)
except Exception as e:
click.echo('\n'.join([
'Oops! Something went wrong while creating the database for you.',
'You may continue with the quicksetup, however:',
'For aiida to work correctly you will have to do that yourself as follows.',
manual_setup_instructions(dbuser=dbuser, dbname=dbname), '',
'Or setup your (OS-level) user to have permissions to create databases and rerun quicksetup.',
''
]))
raise e
# create a profile, by default 'quicksetup' and prompt the user if
# already exists
confs = get_or_create_config()
profile_name = profile or 'quicksetup'
write_profile = False
while not write_profile:
if profile_name in confs.get('profiles', {}):
if click.confirm(
'overwrite existing profile {}?'.format(profile_name)):
write_profile = True
else:
profile_name = click.prompt('new profile name', type=str)
else:
write_profile = True
dbhost = postgres.dbinfo.get('host', 'localhost')
dbport = postgres.dbinfo.get('port', '5432')
from os.path import isabs
repo = repo or 'repository-{}/'.format(profile_name)
if not isabs(repo):
repo = os.path.join(aiida_dir, repo)
setup_args = {
'backend': backend,
'email': email,
'db_host': dbhost,
'db_port': dbport,
'db_name': dbname,
'db_user': dbuser,
'db_pass': dbpass,
'repo': repo,
'first_name': first_name,
'last_name': last_name,
'institution': institution,
'force_overwrite': write_profile,
}
setup(profile_name, only_config=False, non_interactive=True, **setup_args)
# Loop over all valid processes and check if a default profile is set for them
# If not set the newly created profile as default, otherwise prompt whether to override
from aiida.cmdline.commands.profile import valid_processes
default_profiles = confs.get('default_profiles', {})
for process in valid_processes:
# if the user specifies whether to override that's fine
if set_default in [True, False]:
_set_default = set_default
# otherwise we may need to ask
else:
default_profile = default_profiles.get(process, '')
if default_profile:
_set_default = click.confirm(
"The default profile for the '{}' process is set to '{}': "
"do you want to set the newly created '{}' as the new default? (can be reverted later)"
.format(process, default_profile, profile_name))
# if there are no other default profiles, we don't need to ask
else:
_set_default = True
if _set_default:
set_default_profile(process, profile_name, force_rewrite=True)
def quicksetup(profile_name, only_config, set_default, non_interactive, backend, db_host, db_port, db_name, db_username,
db_password, repository, email, first_name, last_name, institution):
"""Set up a sane configuration with as little interaction as possible."""
from aiida.common.setup import create_base_dirs, AIIDA_CONFIG_FOLDER
create_base_dirs()
aiida_dir = os.path.expanduser(AIIDA_CONFIG_FOLDER)
# access postgres
postgres = Postgres(host=db_host, port=db_port, interactive=bool(not non_interactive), quiet=False)
postgres.set_setup_fail_callback(prompt_db_info)
success = postgres.determine_setup()
if not success:
sys.exit(1)
# default database name is <profile_name>_<login-name>
# this ensures that for profiles named test_... the database will also
# be named test_...
import getpass
osuser = getpass.getuser()
dbname = db_name or profile_name + '_' + osuser
# default database user name is aiida_qs_<login-name>
# default password is random
dbuser = db_username or 'aiida_qs_' + osuser
from aiida.common.setup import generate_random_secret_key
dbpass = db_password or generate_random_secret_key()
# check if there is a profile that contains the db user already
# and if yes, take the db user password from there
# This is ok because a user can only see his own config files
from aiida.common.setup import get_or_create_config
confs = get_or_create_config()
profs = confs.get('profiles', {})
for profile in profs.itervalues():
if profile.get('AIIDADB_USER', '') == dbuser and not db_password:
dbpass = profile.get('AIIDADB_PASS')
print 'using found password for {}'.format(dbuser)
break
try:
create = True
if not postgres.dbuser_exists(dbuser):
postgres.create_dbuser(dbuser, dbpass)
else:
dbname, create = _check_db_name(dbname, postgres)
if create:
postgres.create_db(dbuser, dbname)
except Exception as exception:
click.echo('\n'.join([
'Oops! Something went wrong while creating the database for you.',
'You may continue with the quicksetup, however:',
'For aiida to work correctly you will have to do that yourself as follows.',
manual_setup_instructions(dbuser=dbuser, dbname=dbname), '',
'Or setup your (OS-level) user to have permissions to create databases and rerun quicksetup.', ''
]))
raise exception
# create a profile, by default 'quicksetup' and prompt the user if
# already exists
confs = get_or_create_config()
profile_name = profile_name or 'quicksetup'
write_profile = False
while not write_profile:
if profile_name in confs.get('profiles', {}):
if click.confirm('overwrite existing profile {}?'.format(profile_name)):
write_profile = True
else:
profile_name = click.prompt('new profile name', type=str)
else:
write_profile = True
dbhost = postgres.dbinfo.get('host', 'localhost')
dbport = postgres.dbinfo.get('port', '5432')
from os.path import isabs
repo = repository or 'repository-{}/'.format(profile_name)
if not isabs(repo):
repo = os.path.join(aiida_dir, repo)
setup_args = {
'backend': backend,
'email': email,
'db_host': dbhost,
'db_port': dbport,
'db_name': dbname,
'db_user': dbuser,
'db_pass': dbpass,
'repo': repo,
'first_name': first_name,
'last_name': last_name,
'institution': institution,
'force_overwrite': write_profile,
}
setup_profile(profile_name, only_config=only_config, set_default=set_default, non_interactive=True, **setup_args)