Преглед на файлове

Reduce duplication with a common function

Edward Shornock преди 5 години
родител
ревизия
33113890f5
променени са 1 файла, в които са добавени 21 реда и са изтрити 25 реда
  1. 21 25
      borgmatic/hooks/postgresql.py

+ 21 - 25
borgmatic/hooks/postgresql.py

@@ -15,6 +15,25 @@ def make_dump_path(location_config):  # pragma: no cover
     )
 
 
+def make_extra_environment(database):
+    '''
+    Make the extra_environment dict from the given database configuration.
+    '''
+    extra = dict()
+    if 'password' in database:
+        extra['PGPASSWORD'] = database['password']
+    extra['PGSSLMODE'] = database.get('sslmode', 'disable')
+    if 'sslcert' in database:
+        extra['PGSSLCERT'] = database['sslcert']
+    if 'sslkey' in database:
+        extra['PGSSLKEY'] = database['sslkey']
+    if 'sslrootcert' in database:
+        extra['PGSSLROOTCERT'] = database['sslrootcert']
+    if 'sslcrl' in database:
+        extra['PGSSLCRL'] = database['sslcrl']
+    return extra
+
+
 def dump_databases(databases, log_prefix, location_config, dry_run):
     '''
     Dump the given PostgreSQL databases to a named pipe. The databases are supplied as a sequence of
@@ -56,19 +75,7 @@ def dump_databases(databases, log_prefix, location_config, dry_run):
             # format in a particular, a named destination is required, and redirection doesn't work.
             + (('>', dump_filename) if dump_format != 'directory' else ())
         )
-        extra_environment = dict()
-        if 'password' in database:
-            extra_environment['PGPASSWORD'] = database['password']
-        extra_environment['PGSSLMODE'] = database['sslmode'] if 'sslmode' in database else 'disable'
-        if 'sslcert' in database:
-            extra_environment['PGSSLCERT'] = database['sslcert']
-        if 'sslkey' in database:
-            extra_environment['PGSSLKEY'] = database['sslkey']
-        if 'sslrootcert' in database:
-            extra_environment['PGSSLROOTCERT'] = database['sslrootcert']
-        if 'sslcrl' in database:
-            extra_environment['PGSSLCRL'] = database['sslcrl']
-
+        extra_environment = make_extra_environment(database)
 
         logger.debug(
             '{}: Dumping PostgreSQL database {} to {}{}'.format(
@@ -153,18 +160,7 @@ def restore_database_dump(database_config, log_prefix, location_config, dry_run,
         + (('--username', database['username']) if 'username' in database else ())
         + (() if extract_process else (dump_filename,))
     )
-    extra_environment = dict()
-    if 'password' in database:
-        extra_environment['PGPASSWORD'] = database['password']
-    extra_environment['PGSSLMODE'] = database['sslmode'] if 'sslmode' in database else 'disable'
-    if 'sslcert' in database:
-        extra_environment['PGSSLCERT'] = database['sslcert']
-    if 'sslkey' in database:
-        extra_environment['PGSSLKEY'] = database['sslkey']
-    if 'sslrootcert' in database:
-        extra_environment['PGSSLROOTCERT'] = database['sslrootcert']
-    if 'sslcrl' in database:
-        extra_environment['PGSSLCRL'] = database['sslcrl']
+    extra_environment = make_extra_environment(database)
 
     logger.debug(
         '{}: Restoring PostgreSQL database {}{}'.format(log_prefix, database['name'], dry_run_label)