From 114dec13a16e7ebef9aa0eb09468a44a3c795110 Mon Sep 17 00:00:00 2001 From: Dominik George Date: Wed, 17 Feb 2021 23:33:49 +0100 Subject: [PATCH] Use environment and pgpass to connect to PostgreSQL Closes #384. --- dbbackup/db/postgresql.py | 89 ++++++++++++++++++++++++--------------- 1 file changed, 56 insertions(+), 33 deletions(-) diff --git a/dbbackup/db/postgresql.py b/dbbackup/db/postgresql.py index 89707ca0..a1a2971c 100644 --- a/dbbackup/db/postgresql.py +++ b/dbbackup/db/postgresql.py @@ -1,5 +1,6 @@ -from urllib.parse import quote +from tempfile import mkstemp import logging +import os from .base import BaseCommandDBConnector from .exceptions import DumpError @@ -7,23 +8,48 @@ logger = logging.getLogger('dbbackup.command') -def create_postgres_uri(self): - host = self.settings.get('HOST') - if not host: - raise DumpError('A host name is required') - - dbname = self.settings.get('NAME') or '' - user = quote(self.settings.get('USER') or '') - password = self.settings.get('PASSWORD') or '' - password = ':{}'.format(quote(password)) if password else '' - if not user: - password = '' - else: - host = '@' + host - - port = ':{}'.format(self.settings.get('PORT')) if self.settings.get('PORT') else '' - dbname = f'--dbname=postgresql://{user}{password}{host}{port}/{dbname}' - return dbname +class PgEnvWrapper: + """ + Context manager that updates the OS environment with the libpq variables + derived from settings, and if necessary a temporary .pgpass file. + """ + def __init__(self, settings): + self.settings = settings + self.pgpass_fd, self.pgpass_path = None, None + + def __enter__(self): + # Get all settings, with empty defaults to detect later + pghost = self.settings.get('HOST', None) + pgport = self.settings.get('PORT', None) + pguser = self.settings.get('USER', None) + pgdatabase = self.settings.get('NAME', None) + pgpassword = self.settings.get('PASSWORD', None) + + # Set PG* environment variables for everything we got + # All defaults are thus left to libpq + env = os.environ.copy() + if pghost: + env['PGHOST'] = pghost + if pgport: + env['PGPORT'] = pgport + if pguser: + env['PGUSER'] = pguser + if pgdatabase: + env['PGDATABASE'] = pgdatabase + + if pgpassword: + # Open a temporary file (safe name, mode 600) as .pgpass file + fd, self.pgpass_path = mkstemp(text=True) + os.close(fd) + with open(self.pgpass_path, 'w') as pgpass_file: + # Write a catch-all entry, as this .pgass is only used once and by us + pgpass_file.write(f'*:*:*:*:{pgpassword}\n') + env['PGPASSFILE'] = self.pgpass_path + + return env + + def __exit__(self, *args): + os.unlink(self.pgpass_path) class PgDumpConnector(BaseCommandDBConnector): @@ -39,7 +65,6 @@ class PgDumpConnector(BaseCommandDBConnector): def _create_dump(self): cmd = '{} '.format(self.dump_cmd) - cmd = cmd + create_postgres_uri(self) for table in self.exclude: cmd += ' --exclude-table-data={}'.format(table) @@ -47,20 +72,20 @@ def _create_dump(self): cmd += ' --clean' cmd = '{} {} {}'.format(self.dump_prefix, cmd, self.dump_suffix) - stdout, stderr = self.run_command(cmd, env=self.dump_env) + with PgEnvWrapper(self.settings) as env: + stdout, stderr = self.run_command(cmd, env={**self.dump_env, **env}) return stdout def _restore_dump(self, dump): cmd = '{} '.format(self.restore_cmd) - cmd = cmd + create_postgres_uri(self) # without this, psql terminates with an exit value of 0 regardless of errors cmd += ' --set ON_ERROR_STOP=on' if self.single_transaction: cmd += ' --single-transaction' - cmd += ' {}'.format(self.settings['NAME']) cmd = '{} {} {}'.format(self.restore_prefix, cmd, self.restore_suffix) - stdout, stderr = self.run_command(cmd, stdin=dump, env=self.restore_env) + with PgEnvWrapper(self.settings) as env: + stdout, stderr = self.run_command(cmd, stdin=dump, env={**self.restore_env, **env}) return stdout, stderr @@ -76,11 +101,8 @@ def _enable_postgis(self): self.psql_cmd) cmd += ' --username={}'.format(self.settings['ADMIN_USER']) cmd += ' --no-password' - if self.settings.get('HOST'): - cmd += ' --host={}'.format(self.settings['HOST']) - if self.settings.get('PORT'): - cmd += ' --port={}'.format(self.settings['PORT']) - return self.run_command(cmd) + with PgEnvWrapper(self.settings) as env: + return self.run_command(cmd, env=env) def _restore_dump(self, dump): if self.settings.get('ADMIN_USER'): @@ -101,23 +123,24 @@ class PgDumpBinaryConnector(PgDumpConnector): def _create_dump(self): cmd = '{} '.format(self.dump_cmd) - cmd = cmd + create_postgres_uri(self) cmd += ' --format=custom' for table in self.exclude: cmd += ' --exclude-table-data={}'.format(table) cmd = '{} {} {}'.format(self.dump_prefix, cmd, self.dump_suffix) - stdout, stderr = self.run_command(cmd, env=self.dump_env) + with PgEnvWrapper(self.settings) as env: + stdout, stderr = self.run_command(cmd, env={**self.dump_env, **env}) return stdout def _restore_dump(self, dump): - dbname = create_postgres_uri(self) - cmd = '{} {}'.format(self.restore_cmd, dbname) + cmd = '{} '.format(self.restore_cmd) if self.single_transaction: cmd += ' --single-transaction' if self.drop: cmd += ' --clean' + cmd += '-d {}'.format(self.settings.get('NAME')) cmd = '{} {} {}'.format(self.restore_prefix, cmd, self.restore_suffix) - stdout, stderr = self.run_command(cmd, stdin=dump, env=self.restore_env) + with PgEnvWrapper(self.settings) as env: + stdout, stderr = self.run_command(cmd, stdin=dump, env={**self.restore_env, **env}) return stdout, stderr