Use environment and pgpass to connect to PostgreSQL

Closes jazzband#384.

Author: Natureshadow

Diff for: dbbackup/db/postgresql.py

@@ -1,29 +1,46 @@
-from urllib.parse import quote
+from tempfile import mkstemp
 import logging
+import os
 
 from .base import BaseCommandDBConnector
 from .exceptions import DumpError
 
 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')
+class PgEnvWrapper:
+    def __init__(self, settings):
+        self.settings = settings
 
-    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
+    def __enter__(self):
+        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)
 
-    port = ':{}'.format(self.settings.get('PORT')) if self.settings.get('PORT') else ''
-    dbname = f'--dbname=postgresql://{user}{password}{host}{port}/{dbname}'
-    return dbname
+        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:
+            self.pgpass_fd, self.pgpass_path = mkstemp(text=True)
+            os.write(self.pgpass_fd, f'*:*:*:*:{pgpassword}\n')
+            os.close(self.pgpass_fd)
+            env['PGPASSFILE'] = self.pgpass_path
+        else:
+            self.pgpass_fd, self.pgpass_path = None, None
+
+        return env
+
+    def __exit__(self):
+        os.unlink(self.pgpass_path)
 
 
 class PgDumpConnector(BaseCommandDBConnector):
@@ -39,28 +56,27 @@ 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)
         if self.drop:
             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 +92,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 +114,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