From 2ed891401bc097a672130285adc1d6bc22063afd Mon Sep 17 00:00:00 2001
From: Dominik George <nik@naturalnet.de>
Date: Wed, 17 Feb 2021 22:32:04 +0100
Subject: [PATCH] Use PG* environment rather than command-line for PostgreSQL

Closes #384.
---
 dbbackup/db/postgresql.py | 38 ++++++++++++++------------------------
 1 file changed, 14 insertions(+), 24 deletions(-)

diff --git a/dbbackup/db/postgresql.py b/dbbackup/db/postgresql.py
index 89707ca0..1ba5518b 100644
--- a/dbbackup/db/postgresql.py
+++ b/dbbackup/db/postgresql.py
@@ -1,5 +1,6 @@
 from urllib.parse import quote
 import logging
+import os
 
 from .base import BaseCommandDBConnector
 from .exceptions import DumpError
@@ -7,25 +8,6 @@
 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 PgDumpConnector(BaseCommandDBConnector):
     """
     PostgreSQL connector, it uses pg_dump`` to create an SQL text file
@@ -37,9 +19,20 @@ class PgDumpConnector(BaseCommandDBConnector):
     single_transaction = True
     drop = True
 
+    @property
+    def dump_env(self):
+        env = os.environ.copy()
+        for var in ('HOST', 'PORT', 'USER', 'PASSWORD'):
+            if self.settings.get(var, None):
+                env['PG%s' % var] = self.settings.get(var)
+        env['PGDATABASE'] = self.settings.get(var, 'postgres')
+
+        return env
+
+    restore_env = dump_env
+
     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)
@@ -52,7 +45,6 @@ def _create_dump(self):
 
     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'
@@ -101,7 +93,6 @@ 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:
@@ -111,8 +102,7 @@ def _create_dump(self):
         return stdout
 
     def _restore_dump(self, dump):
-        dbname = create_postgres_uri(self)
-        cmd = '{} {}'.format(self.restore_cmd, dbname)
+        cmd = '{} {}'.format(self.restore_cmd, self.settings['NAME'])
 
         if self.single_transaction:
             cmd += ' --single-transaction'