Skip to content

Commit

Permalink
Move aiida.manage.external.pgsu to external package pgsu (#3892)
Browse files Browse the repository at this point in the history 
One of the main issues with `verdi quicksetup` is that it is non-trivial
to figure out how to connect as the PostgreSQL superuser in a wide
variety of operating systems and PostgreSQL setups. This PR factors out
the code which takes care of this connection into a separate package, 
called `pgsu` which can then be tested on a wide variety of setups using
continuous integration.
  • Loading branch information
@ltalirz
ltalirz authored Apr 6, 2020
1 parent 5c0d86f commit 5da7120
Show file tree
Hide file tree
Showing 13 changed files with 141 additions and 421 deletions.
21 changes: 7 additions & 14 deletions aiida/cmdline/commands/cmd_setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand Down
349 changes: 8 additions & 341 deletions aiida/manage/external/pgsu.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Loading

0 comments on commit 5da7120

Please sign in to comment.