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)