Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

move PGSU out to separate package #3892

Merged
merged 3 commits into from
Apr 6, 2020
Merged

move PGSU out to separate package #3892

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ffbe6a2
move PGSU out to separate package
ltalirz Apr 4, 2020
dc98acd
update requirements files manually
ltalirz Apr 5, 2020
a7c4ad5
fix docs
ltalirz Apr 5, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 7 additions & 14 deletions aiida/cmdline/commands/cmd_setup.py
View file
Open in desktop
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
Open in desktop
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