diff --git a/aiida/cmdline/commands/cmd_setup.py b/aiida/cmdline/commands/cmd_setup.py index 4edc23426b..fbfbf8b23c 100644 --- a/aiida/cmdline/commands/cmd_setup.py +++ b/aiida/cmdline/commands/cmd_setup.py @@ -136,22 +136,15 @@ def quicksetup( echo.echo_critical('failed to determine the PostgreSQL setup') try: - create = True - if not postgres.dbuser_exists(db_username): - postgres.create_dbuser(db_username, db_password) - else: - db_name, create = postgres.check_db_name(db_name) - - if create: - postgres.create_db(db_username, db_name) + db_username, db_name = postgres.create_dbuser_db_safe(dbname=db_name, dbuser=db_username, dbpass=db_password) except Exception as exception: echo.echo_error( '\n'.join([ 'Oops! quicksetup was unable to create the AiiDA database for you.', - 'For AiiDA to work, please either create the database yourself as follows:', - manual_setup_instructions(dbuser=su_db_username, dbname=su_db_name), '', - 'Alternatively, give your (operating system) user permission to create postgresql databases' + - 'and run quicksetup again.', '' + 'See `verdi quicksetup -h` for how to specify non-standard parameters for the postgresql connection.\n' + 'Alternatively, create the AiiDA database yourself: ', + manual_setup_instructions(dbuser=su_db_username, + dbname=su_db_name), '', 'and then use `verdi setup` instead', '' ]) ) raise exception @@ -169,8 +162,8 @@ def quicksetup( 'db_backend': db_backend, 'db_name': db_name, # from now on we connect as the AiiDA DB user, which may be forbidden when going via sockets - 'db_host': db_host or 'localhost', - 'db_port': db_port, + 'db_host': postgres.host_for_psycopg2, + 'db_port': postgres.port_for_psycopg2, 'db_username': db_username, 'db_password': db_password, 'repository': repository, diff --git a/aiida/manage/external/pgsu.py b/aiida/manage/external/pgsu.py index 025972e8a3..05c58e2a97 100644 --- a/aiida/manage/external/pgsu.py +++ b/aiida/manage/external/pgsu.py @@ -13,344 +13,11 @@ separate package that can then be tested on multiple OS / postgres setups. Therefore, **please keep this module entirely AiiDA-agnostic**. """ - -try: - import subprocess32 as subprocess -except ImportError: - import subprocess - -from enum import IntEnum -import click - -DEFAULT_DBINFO = { - 'host': 'localhost', - 'port': 5432, - 'user': 'postgres', - 'password': None, - 'database': 'template1', -} - - -class PostgresConnectionMode(IntEnum): - """Describe mode of connecting to postgres.""" - - DISCONNECTED = 0 - PSYCOPG = 1 - PSQL = 2 - - -class PGSU: - """ - Connect to an existing PostgreSQL cluster as the `postgres` superuser and execute SQL commands. - - Tries to use psycopg2 with a fallback to psql subcommands (using ``sudo su`` to run as postgres user). - - Simple Example:: - - postgres = PGSU() - postgres.execute("CREATE USER testuser PASSWORD 'testpw'") - - Complex Example:: - - postgres = PGSU(interactive=True, dbinfo={'port': 5433}) - postgres.execute("CREATE USER testuser PASSWORD 'testpw'") - - Note: In postgresql - * you cannot drop databases you are currently connected to - * 'template0' is the unmodifiable template database (which you cannot connect to) - * 'template1' is the modifiable template database (which you can connect to) - """ - - def __init__(self, interactive=False, quiet=True, dbinfo=None, determine_setup=True): - """Store postgres connection info. - - :param interactive: use True for verdi commands - :param quiet: use False to show warnings/exceptions - :param dbinfo: psycopg dictionary containing keys like 'host', 'user', 'port', 'database' - :param determine_setup: Whether to determine setup upon instantiation. - You may set this to False and use the 'determine_setup()' method instead. - """ - self.interactive = interactive - self.quiet = quiet - self.connection_mode = PostgresConnectionMode.DISCONNECTED - - self.setup_fail_callback = prompt_db_info if interactive else None - self.setup_fail_counter = 0 - self.setup_max_tries = 1 - - self.dbinfo = DEFAULT_DBINFO.copy() - if dbinfo is not None: - self.dbinfo.update(dbinfo) - - if determine_setup: - self.determine_setup() - - def execute(self, command, **kwargs): - """Execute postgres command using determined connection mode. - - :param command: A psql command line as a str - :param kwargs: will be forwarded to _execute_... function - """ - # Use self.dbinfo as default kwargs, update with provided kwargs - kw_copy = self.dbinfo.copy() - kw_copy.update(kwargs) - - if self.connection_mode == PostgresConnectionMode.PSYCOPG: # pylint: disable=no-else-return - return _execute_psyco(command, **kw_copy) - elif self.connection_mode == PostgresConnectionMode.PSQL: - return _execute_psql(command, **kw_copy) - - raise ValueError('Could not connect to postgres.') - - def set_setup_fail_callback(self, callback): - """ - Set a callback to be called when setup cannot be determined automatically - - :param callback: a callable with signature ``callback(interactive, dbinfo)`` - that returns a ``dbinfo`` dictionary. - """ - self.setup_fail_callback = callback - - def determine_setup(self): - """Determine how to connect as the postgres superuser. - - Depending on how postgres is set up, psycopg2 can be used to create dbs and db users, - otherwise a subprocess has to be used that executes psql as an os user with appropriate permissions. - - Note: We aim to connect as a superuser (typically 'postgres') with privileges to manipulate (create/drop) - databases and database users. - - :returns success: True, if connection could be established. - :rtype success: bool - """ - # find out if we run as a postgres superuser or can connect as postgres - # This will work on OSX in some setups but not in the default Debian one - dbinfo = self.dbinfo.copy() - - for pg_user in set([dbinfo.get('user'), None]): - dbinfo['user'] = pg_user - if _try_connect_psycopg(**dbinfo): - self.dbinfo = dbinfo - self.connection_mode = PostgresConnectionMode.PSYCOPG - return True - - # This will work for the default Debian postgres setup, assuming that sudo is available to the user - # Check if the user can find the sudo command - if _sudo_exists(): - if _try_subcmd(interactive=self.interactive, quiet=self.quiet, **dbinfo): - self.dbinfo = dbinfo - self.connection_mode = PostgresConnectionMode.PSQL - return True - elif not self.quiet: - click.echo('Warning: Could not find `sudo` for connecting to the database.') - - self.setup_fail_counter += 1 - return self._no_setup_detected() - - def _no_setup_detected(self): - """Print a warning message and calls the failed setup callback - - :returns: False, if no successful try. - """ - message = '\n'.join([ - 'Warning: Unable to autodetect postgres setup - do you know how to access it?', - ]) - - if not self.quiet: - click.echo(message) - - if self.setup_fail_callback and self.setup_fail_counter <= self.setup_max_tries: - self.dbinfo = self.setup_fail_callback(self.interactive, self.dbinfo) - return self.determine_setup() - - return False - - @property - def is_connected(self): - return self.connection_mode in (PostgresConnectionMode.PSYCOPG, PostgresConnectionMode.PSQL) - - -def prompt_db_info(interactive, dbinfo): - """ - Prompt interactively for postgres database connection details - - Can be used as a setup fail callback for :py:class:`PGSU` - - :return: dictionary with the following keys: host, port, database, user - """ - if not interactive: - return DEFAULT_DBINFO - - access = False - while not access: - dbinfo_new = {} - dbinfo_new['host'] = click.prompt('postgres host', default=dbinfo.get('host'), type=str) - dbinfo_new['port'] = click.prompt('postgres port', default=dbinfo.get('port'), type=int) - dbinfo_new['user'] = click.prompt('postgres super user', default=dbinfo.get('user'), type=str) - dbinfo_new['database'] = click.prompt('database', default=dbinfo.get('database'), type=str) - click.echo('') - click.echo('Trying to access postgres ...') - if _try_connect_psycopg(**dbinfo_new): - access = True - else: - dbinfo_new['password'] = click.prompt( - 'postgres password of {}'.format(dbinfo_new['user']), hide_input=True, type=str, default='' - ) - if not dbinfo_new.get('password'): - dbinfo_new.pop('password') - return dbinfo_new - - -def _try_connect_psycopg(**kwargs): - """ - try to start a psycopg2 connection. - - :return: True if successful, False otherwise - """ - from psycopg2 import connect - success = False - try: - conn = connect(**kwargs) - success = True - conn.close() - except Exception: # pylint: disable=broad-except - pass - return success - - -def _sudo_exists(): - """ - Check that the sudo command can be found - - :return: True if successful, False otherwise - """ - try: - subprocess.check_output(['sudo', '-V']) - except subprocess.CalledProcessError: - return False - except OSError: - return False - - return True - - -def _try_subcmd(**kwargs): - """ - try to run psql in a subprocess. - - :return: True if successful, False otherwise - """ - success = False - try: - kwargs['stderr'] = subprocess.STDOUT - _execute_psql(r'\q', **kwargs) - success = True - except subprocess.CalledProcessError: - pass - return success - - -def _execute_psyco(command, **kwargs): - """ - executes a postgres commandline through psycopg2 - - :param command: A psql command line as a str - :param kwargs: will be forwarded to psycopg2.connect - """ - import psycopg2 - - # Note: Ubuntu 18.04 uses "peer" as the default postgres configuration - # which allows connections only when the unix user matches the database user. - # This restriction no longer applies for IPv4/v6-based connection, - # when specifying host=localhost. - if kwargs.get('host') is None: - kwargs['host'] = 'localhost' - - output = None - with psycopg2.connect(**kwargs) as conn: - conn.autocommit = True - with conn.cursor() as cursor: - cursor.execute(command) - if cursor.description is not None: - output = cursor.fetchall() - - # see http://initd.org/psycopg/docs/usage.html#with-statement - conn.close() - return output - - -def _execute_psql(command, user='postgres', quiet=True, interactive=False, **kwargs): - """ - Executes an SQL command via ``psql`` as another system user in a subprocess. - - Tries to "become" the user specified in ``kwargs`` (i.e. interpreted as UNIX system user) - and run psql in a subprocess. - - :param command: A psql command line as a str - :param quiet: If True, don't print warnings. - :param interactive: If False, `sudo` won't ask for a password and fail if one is required. - :param kwargs: connection details to forward to psql, signature as in psycopg2.connect - """ - option_str = '' - - database = kwargs.pop('database', None) - if database: - option_str += '-d {}'.format(database) - # to do: Forward password to psql; ignore host only when the password is None. # pylint: disable=fixme - kwargs.pop('password', None) - - host = kwargs.pop('host', 'localhost') - if host and host != 'localhost': - option_str += ' -h {}'.format(host) - elif not quiet: - click.echo( - "Warning: Found host 'localhost' but dropping '-h localhost' option for psql " + - 'since this may cause psql to switch to password-based authentication.' - ) - - port = kwargs.pop('port', None) - if port: - option_str += ' -p {}'.format(port) - - user = kwargs.pop('user', 'postgres') - - # Build command line - sudo_cmd = ['sudo'] - if not interactive: - sudo_cmd += ['-n'] - su_cmd = ['su', user, '-c'] - - psql_cmd = ['psql {opt} -tc {cmd}'.format(cmd=escape_for_bash(command), opt=option_str)] - sudo_su_psql = sudo_cmd + su_cmd + psql_cmd - result = subprocess.check_output(sudo_su_psql, **kwargs) - result = result.decode('utf-8').strip().split('\n') - result = [i for i in result if i] - - return result - - -def escape_for_bash(str_to_escape): - """ - This function takes any string and escapes it in a way that - bash will interpret it as a single string. - - Explanation: - - At the end, in the return statement, the string is put within single - quotes. Therefore, the only thing that I have to escape in bash is the - single quote character. To do this, I substitute every single - quote ' with '"'"' which means: - - First single quote: exit from the enclosing single quotes - - Second, third and fourth character: "'" is a single quote character, - escaped by double quotes - - Last single quote: reopen the single quote to continue the string - - Finally, note that for python I have to enclose the string '"'"' - within triple quotes to make it work, getting finally: the complicated - string found below. - """ - escaped_quotes = str_to_escape.replace("'", """'"'"'""") - return "'{}'".format(escaped_quotes) +import warnings +from pgsu import PGSU, PostgresConnectionMode, DEFAULT_DSN as DEFAULT_DBINFO # pylint: disable=unused-import,no-name-in-module +from aiida.common.warnings import AiidaDeprecationWarning + +warnings.warn( # pylint: disable=no-member + '`aiida.manage.external.pgsu` is now available in the separate `pgsu` package. ' + 'This module will be removed entirely in AiiDA 2.0.0', AiidaDeprecationWarning +) diff --git a/aiida/manage/external/postgres.py b/aiida/manage/external/postgres.py index 0c6e9b1c5a..680b62e088 100644 --- a/aiida/manage/external/postgres.py +++ b/aiida/manage/external/postgres.py @@ -21,7 +21,7 @@ import click from aiida.cmdline.utils import echo -from .pgsu import PGSU, PostgresConnectionMode, DEFAULT_DBINFO +from pgsu import PGSU, PostgresConnectionMode, DEFAULT_DSN as DEFAULT_DBINFO # pylint: disable=no-name-in-module _CREATE_USER_COMMAND = 'CREATE USER "{}" WITH PASSWORD \'{}\'' _DROP_USER_COMMAND = 'DROP USER "{}"' @@ -32,20 +32,20 @@ ) _DROP_DB_COMMAND = 'DROP DATABASE "{}"' _GRANT_PRIV_COMMAND = 'GRANT ALL PRIVILEGES ON DATABASE "{}" TO "{}"' -_GET_USERS_COMMAND = "SELECT usename FROM pg_user WHERE usename='{}'" +_USER_EXISTS_COMMAND = "SELECT usename FROM pg_user WHERE usename='{}'" _CHECK_DB_EXISTS_COMMAND = "SELECT datname FROM pg_database WHERE datname='{}'" _COPY_DB_COMMAND = 'CREATE DATABASE "{}" WITH TEMPLATE "{}" OWNER "{}"' class Postgres(PGSU): """ - Adds convenience functions to pgsu.Postgres. + Adds convenience functions to :py:class:`pgsu.PGSU`. - Provides conenience functions for + Provides convenience functions for * creating/dropping users * creating/dropping databases - etc. See pgsu.Postgres for implementation details. + etc. Example:: @@ -55,6 +55,10 @@ class Postgres(PGSU): postgres.create_db('username', 'dbname') """ + def __init__(self, dbinfo=None, **kwargs): + """See documentation of :py:meth:`pgsu.PGSU.__init__`.""" + super().__init__(dsn=dbinfo, **kwargs) + @classmethod def from_profile(cls, profile, **kwargs): """Create Postgres instance with dbinfo from AiiDA profile data. @@ -63,7 +67,7 @@ def from_profile(cls, profile, **kwargs): database superuser. :param profile: AiiDA profile instance - :param kwargs: keyword arguments forwarded to Postgres constructor + :param kwargs: keyword arguments forwarded to PGSU constructor :returns: Postgres instance pre-populated with data from AiiDA profile """ @@ -77,23 +81,25 @@ def from_profile(cls, profile, **kwargs): return Postgres(dbinfo=dbinfo, **kwargs) - def check_db_name(self, dbname): - """Looks up if a database with the name exists, prompts for using or creating a differently named one.""" - create = True - while create and self.db_exists(dbname): - echo.echo_info('database {} already exists!'.format(dbname)) - if not click.confirm('Use it (make sure it is not used by another profile)?'): - dbname = click.prompt('new name', type=str, default=dbname) - else: - create = False - return dbname, create + ### DB user functions ### + + def dbuser_exists(self, dbuser): + """ + Find out if postgres user with name dbuser exists + + :param str dbuser: database user to check for + :return: (bool) True if user exists, False otherwise + """ + return bool(self.execute(_USER_EXISTS_COMMAND.format(dbuser))) def create_dbuser(self, dbuser, dbpass): """ Create a database user in postgres - :param dbuser: (str), Name of the user to be created. - :param dbpass: (str), Password the user should be given. + :param str dbuser: Name of the user to be created. + :param str dbpass: Password the user should be given. + :raises: psycopg2.errors.DuplicateObject if user already exists and + self.connection_mode == PostgresConnectionMode.PSYCOPG """ self.execute(_CREATE_USER_COMMAND.format(dbuser, dbpass)) @@ -101,25 +107,42 @@ def drop_dbuser(self, dbuser): """ Drop a database user in postgres - :param dbuser: (str), Name of the user to be dropped. + :param str dbuser: Name of the user to be dropped. """ self.execute(_DROP_USER_COMMAND.format(dbuser)) - def dbuser_exists(self, dbuser): + def check_dbuser(self, dbuser): + """Looks up if a given user already exists, prompts for using or creating a differently named one. + + :param str dbuser: Name of the user to be created or reused. + :returns: tuple (dbuser, created) """ - Find out if postgres user with name dbuser exists + create = True + while create and self.dbuser_exists(dbuser): + echo.echo_info('Database user "{}" already exists!'.format(dbuser)) + if not click.confirm('Use it? '): + dbuser = click.prompt('New database user name: ', type=str, default=dbuser) + else: + create = False + return dbuser, create - :param dbuser: (str) database user to check for - :return: (bool) True if user exists, False otherwise + ### DB functions ### + + def db_exists(self, dbname): """ - return bool(self.execute(_GET_USERS_COMMAND.format(dbuser))) + Check wether a postgres database with dbname exists + + :param str dbname: Name of the database to check for + :return: (bool), True if database exists, False otherwise + """ + return bool(self.execute(_CHECK_DB_EXISTS_COMMAND.format(dbname))) def create_db(self, dbuser, dbname): """ Create a database in postgres - :param dbuser: (str), Name of the user which should own the db. - :param dbname: (str), Name of the database. + :param str dbuser: Name of the user which should own the db. + :param str dbname: Name of the database. """ self.execute(_CREATE_DB_COMMAND.format(dbname, dbuser)) self.execute(_GRANT_PRIV_COMMAND.format(dbname, dbuser)) @@ -128,28 +151,70 @@ def drop_db(self, dbname): """ Drop a database in postgres - :param dbname: (str), Name of the database. + :param str dbname: Name of the database. """ self.execute(_DROP_DB_COMMAND.format(dbname)) def copy_db(self, src_db, dest_db, dbuser): self.execute(_COPY_DB_COMMAND.format(dest_db, src_db, dbuser)) - def db_exists(self, dbname): + def check_db(self, dbname): + """Looks up if a database with the name exists, prompts for using or creating a differently named one. + + :param str dbname: Name of the database to be created or reused. + :returns: tuple (dbname, created) """ - Check wether a postgres database with dbname exists + create = True + while create and self.db_exists(dbname): + echo.echo_info('database {} already exists!'.format(dbname)) + if not click.confirm('Use it (make sure it is not used by another profile)?'): + dbname = click.prompt('new name', type=str, default=dbname) + else: + create = False + return dbname, create - :param dbname: Name of the database to check for - :return: (bool), True if database exists, False otherwise + def create_dbuser_db_safe(self, dbname, dbuser, dbpass): + """Create DB and user + grant privileges. + + Prompts when reusing existing users / databases. """ - return bool(self.execute(_CHECK_DB_EXISTS_COMMAND.format(dbname))) + dbuser, create = self.check_dbuser(dbuser=dbuser) + if create: + self.create_dbuser(dbuser=dbuser, dbpass=dbpass) + + dbname, create = self.check_db(dbname=dbname) + if create: + self.create_db(dbuser, dbname) + + return dbuser, dbname + + @property + def host_for_psycopg2(self): + """Return correct host for psycopg2 connection (as required by regular AiiDA operation).""" + host = self.dsn.get('host') + if self.connection_mode == PostgresConnectionMode.PSQL: + # If "sudo su postgres" was needed to create the DB, we are likely on Ubuntu, where + # the same will *not* work for arbitrary database users => enforce TCP/IP connection + host = host or 'localhost' + + return host + + @property + def port_for_psycopg2(self): + """Return port for psycopg2 connection (as required by regular AiiDA operation).""" + return self.dsn.get('port') + + @property + def dbinfo(self): + """Alias for Postgres.dsn.""" + return self.dsn.copy() def manual_setup_instructions(dbuser, dbname): """Create a message with instructions for manually creating a database""" dbpass = '' instructions = '\n'.join([ - 'Please run the following commands as the user for PostgreSQL (Ubuntu: $sudo su postgres):', + 'Run the following commands as a UNIX user with access to PostgreSQL (Ubuntu: $ sudo su postgres):', '', '\t$ psql template1', '\t==> ' + _CREATE_USER_COMMAND.format(dbuser, dbpass), diff --git a/aiida/manage/tests/__init__.py b/aiida/manage/tests/__init__.py index 9ff04e39c8..cc4e2e5fbf 100644 --- a/aiida/manage/tests/__init__.py +++ b/aiida/manage/tests/__init__.py @@ -252,10 +252,11 @@ def __init__(self, backend=BACKEND_DJANGO, pgtest=None): # pylint: disable=supe self.postgres = None self._profile = None self._has_test_db = False - self._backup = {} - self._backup['config'] = configuration.CONFIG - self._backup['config_dir'] = settings.AIIDA_CONFIG_FOLDER - self._backup['profile'] = configuration.PROFILE + self._backup = { + 'config': configuration.CONFIG, + 'config_dir': settings.AIIDA_CONFIG_FOLDER, + 'profile': configuration.PROFILE, + } @property def profile_dictionary(self): @@ -264,10 +265,10 @@ def profile_dictionary(self): Used to set up AiiDA profile from self.profile_info dictionary. """ dictionary = { - 'database_engine': self.profile_info['database_engine'], - 'database_backend': self.profile_info['database_backend'], - 'database_port': self.dbinfo.get('port'), - 'database_hostname': self.dbinfo.get('host'), + 'database_engine': self.profile_info.get('database_engine'), + 'database_backend': self.profile_info.get('database_backend'), + 'database_port': self.profile_info.get('database_port'), + 'database_hostname': self.profile_info.get('database_hostname'), 'database_name': self.profile_info.get('database_name'), 'database_username': self.profile_info.get('database_username'), 'database_password': self.profile_info.get('database_password'), @@ -297,9 +298,12 @@ def create_aiida_db(self): if self.pg_cluster is None: self.create_db_cluster() self.postgres = Postgres(interactive=False, quiet=True, dbinfo=self.dbinfo) - self.dbinfo = self.postgres.dbinfo.copy() + # note: not using postgres.create_dbuser_db_safe here since we don't want prompts self.postgres.create_dbuser(self.profile_info['database_username'], self.profile_info['database_password']) self.postgres.create_db(self.profile_info['database_username'], self.profile_info['database_name']) + self.dbinfo = self.postgres.dbinfo + self.profile_info['database_hostname'] = self.postgres.host_for_psycopg2 + self.profile_info['database_port'] = self.postgres.port_for_psycopg2 self._has_test_db = True def create_profile(self): diff --git a/docs/requirements_for_rtd.txt b/docs/requirements_for_rtd.txt index 4336f2f21c..cb290f8f83 100644 --- a/docs/requirements_for_rtd.txt +++ b/docs/requirements_for_rtd.txt @@ -22,6 +22,7 @@ kiwipy[rmq]~=0.5.1 numpy<1.18,~=1.17 paramiko~=2.6 pg8000~=1.13 +pgsu~=0.1.0 pgtest>=1.3.1,~=1.3 pika~=1.1 plumpy~=0.14.5 diff --git a/docs/source/nitpick-exceptions b/docs/source/nitpick-exceptions index 4019aad3f7..46dc269e12 100644 --- a/docs/source/nitpick-exceptions +++ b/docs/source/nitpick-exceptions @@ -103,3 +103,6 @@ py:class aldjemy.orm.DbLog py:class aldjemy.orm.DbSetting py:class alembic.config.Config + +py:class pgsu.PGSU +py:meth pgsu.PGSU.__init__ diff --git a/environment.yml b/environment.yml index cfbb126c57..d92ed0950a 100644 --- a/environment.yml +++ b/environment.yml @@ -23,6 +23,7 @@ dependencies: - paramiko~=2.6 - pika~=1.1 - plumpy~=0.14.5 +- pgsu~=0.1.0 - psutil~=5.6 - psycopg2>=2.8.3,~=2.8 - python-dateutil~=2.8 diff --git a/requirements/requirements-py-3.5.txt b/requirements/requirements-py-3.5.txt index ff5c36ad78..545e13a27f 100644 --- a/requirements/requirements-py-3.5.txt +++ b/requirements/requirements-py-3.5.txt @@ -74,6 +74,7 @@ pathlib2==2.3.5 pexpect==4.8.0 pg8000==1.13.2 pgtest==1.3.2 +pgsu==0.1.0 pickleshare==0.7.5 pika==1.1.0 pluggy==0.13.1 diff --git a/requirements/requirements-py-3.6.txt b/requirements/requirements-py-3.6.txt index 4e9015bc17..665898918e 100644 --- a/requirements/requirements-py-3.6.txt +++ b/requirements/requirements-py-3.6.txt @@ -73,6 +73,7 @@ paramiko==2.7.1 parso==0.6.2 pexpect==4.8.0 pg8000==1.13.2 +pgsu==0.1.0 pgtest==1.3.2 pickleshare==0.7.5 pika==1.1.0 diff --git a/requirements/requirements-py-3.7.txt b/requirements/requirements-py-3.7.txt index f52fa482bd..3a166e7b53 100644 --- a/requirements/requirements-py-3.7.txt +++ b/requirements/requirements-py-3.7.txt @@ -72,6 +72,7 @@ paramiko==2.7.1 parso==0.6.2 pexpect==4.8.0 pg8000==1.13.2 +pgsu==0.1.0 pgtest==1.3.2 pickleshare==0.7.5 pika==1.1.0 diff --git a/requirements/requirements-py-3.8.txt b/requirements/requirements-py-3.8.txt index 462ddf5777..a3295d91f7 100644 --- a/requirements/requirements-py-3.8.txt +++ b/requirements/requirements-py-3.8.txt @@ -71,6 +71,7 @@ paramiko==2.7.1 parso==0.6.2 pexpect==4.8.0 pg8000==1.13.2 +pgsu==0.1.0 pgtest==1.3.2 pickleshare==0.7.5 pika==1.1.0 diff --git a/setup.json b/setup.json index de5f83b295..ebe8360174 100644 --- a/setup.json +++ b/setup.json @@ -37,6 +37,7 @@ "paramiko~=2.6", "pika~=1.1", "plumpy~=0.14.5", + "pgsu~=0.1.0", "psutil~=5.6", "psycopg2-binary~=2.8,>=2.8.3", "pyblake2~=1.1; python_version<'3.6'", diff --git a/tests/manage/external/test_postgres.py b/tests/manage/external/test_postgres.py index 266dc92921..5b83374f06 100644 --- a/tests/manage/external/test_postgres.py +++ b/tests/manage/external/test_postgres.py @@ -9,16 +9,10 @@ ########################################################################### """Unit tests for postgres database maintenance functionality""" from unittest import TestCase -from unittest.mock import patch from aiida.manage.external.postgres import Postgres -def _try_connect_always_fail(**kwargs): # pylint: disable=unused-argument - """Always return False""" - return False - - class PostgresTest(TestCase): """Test the public API provided by the `Postgres` class""" @@ -38,31 +32,18 @@ def _setup_postgres(self): return Postgres(interactive=False, quiet=True, dbinfo=self.pg_test.dsn) def test_determine_setup_fail(self): + """Check that setup fails, if bad port is provided. + + Note: In interactive mode, this would prompt for the connection details. + """ postgres = Postgres(interactive=False, quiet=True, dbinfo={'port': '11111'}) self.assertFalse(postgres.is_connected) def test_determine_setup_success(self): + """Check that setup works with default parameters.""" postgres = self._setup_postgres() self.assertTrue(postgres.is_connected) - 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 - - postgres = Postgres(interactive=False, quiet=True, dbinfo={'port': '11111'}, determine_setup=False) - postgres.set_setup_fail_callback(correct_setup) - setup_success = postgres.determine_setup() - self.assertTrue(setup_success) - - @patch('aiida.manage.external.pgsu._try_connect_psycopg', new=_try_connect_always_fail) - @patch('aiida.manage.external.pgsu._try_subcmd') - def test_fallback_on_subcmd(self, try_subcmd): - """Ensure that accessing postgres via subcommand is tried if psycopg 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""" postgres = self._setup_postgres()