2
0
Эх сурвалжийг харах

Fix multiple bugs in PostgreSQL hook (#678).

Reviewed-on: https://projects.torsion.org/borgmatic-collective/borgmatic/pulls/677
Dan Helfman 2 жил өмнө
parent
commit
da0f5a34f2

+ 15 - 8
borgmatic/hooks/postgresql.py

@@ -2,6 +2,7 @@ import csv
 import itertools
 import logging
 import os
+import shlex
 
 from borgmatic.execute import (
     execute_command,
@@ -60,8 +61,10 @@ def database_names_to_dump(database, extra_environment, log_prefix, dry_run):
     if dry_run:
         return ()
 
+    psql_command = shlex.split(database.get('psql_command') or 'psql')
     list_command = (
-        ('psql', '--list', '--no-password', '--csv', '--tuples-only')
+        tuple(psql_command)
+        + ('--list', '--no-password', '--no-psqlrc', '--csv', '--tuples-only')
         + (('--host', database['hostname']) if 'hostname' in database else ())
         + (('--port', str(database['port'])) if 'port' in database else ())
         + (('--username', database['username']) if 'username' in database else ())
@@ -210,9 +213,10 @@ def restore_database_dump(database_config, log_prefix, location_config, dry_run,
     dump_filename = dump.make_database_dump_filename(
         make_dump_path(location_config), database['name'], database.get('hostname')
     )
-    psql_command = database.get('psql_command') or 'psql'
+    psql_command = shlex.split(database.get('psql_command') or 'psql')
     analyze_command = (
-        (psql_command, '--no-password', '--quiet')
+        tuple(psql_command)
+        + ('--no-password', '--no-psqlrc', '--quiet')
         + (('--host', database['hostname']) if 'hostname' in database else ())
         + (('--port', str(database['port'])) if 'port' in database else ())
         + (('--username', database['username']) if 'username' in database else ())
@@ -220,14 +224,17 @@ def restore_database_dump(database_config, log_prefix, location_config, dry_run,
         + (tuple(database['analyze_options'].split(' ')) if 'analyze_options' in database else ())
         + ('--command', 'ANALYZE')
     )
-    pg_restore_command = database.get('pg_restore_command') or 'pg_restore'
+    use_psql_command = all_databases or database.get('format') == 'plain'
+    pg_restore_command = shlex.split(database.get('pg_restore_command') or 'pg_restore')
     restore_command = (
-        (psql_command if all_databases else pg_restore_command, '--no-password')
+        tuple(psql_command if use_psql_command else pg_restore_command)
+        + ('--no-password',)
         + (
-            ('--if-exists', '--exit-on-error', '--clean', '--dbname', database['name'])
-            if not all_databases
-            else ()
+            ('--no-psqlrc', '--set', 'ON_ERROR_STOP=on')
+            if use_psql_command
+            else ('--if-exists', '--exit-on-error', '--clean')
         )
+        + (('--dbname', database['name']) if not all_databases else ())
         + (('--host', database['hostname']) if 'hostname' in database else ())
         + (('--port', str(database['port'])) if 'port' in database else ())
         + (('--username', database['username']) if 'username' in database else ())

+ 11 - 7
tests/unit/hooks/test_postgresql.py

@@ -56,6 +56,7 @@ def test_database_names_to_dump_with_all_and_format_lists_databases_with_hostnam
             'psql',
             '--list',
             '--no-password',
+            '--no-psqlrc',
             '--csv',
             '--tuples-only',
             '--host',
@@ -75,7 +76,7 @@ def test_database_names_to_dump_with_all_and_format_lists_databases_with_hostnam
 def test_database_names_to_dump_with_all_and_format_lists_databases_with_username():
     database = {'name': 'all', 'format': 'custom', 'username': 'postgres'}
     flexmock(module).should_receive('execute_command_and_capture_output').with_args(
-        ('psql', '--list', '--no-password', '--csv', '--tuples-only', '--username', 'postgres'),
+        ('psql', '--list', '--no-password', '--no-psqlrc', '--csv', '--tuples-only', '--username', 'postgres'),
         extra_environment=object,
     ).and_return('foo,test,\nbar,test,"stuff and such"')
 
@@ -88,7 +89,7 @@ def test_database_names_to_dump_with_all_and_format_lists_databases_with_usernam
 def test_database_names_to_dump_with_all_and_format_lists_databases_with_options():
     database = {'name': 'all', 'format': 'custom', 'list_options': '--harder'}
     flexmock(module).should_receive('execute_command_and_capture_output').with_args(
-        ('psql', '--list', '--no-password', '--csv', '--tuples-only', '--harder'),
+        ('psql', '--list', '--no-password', '--no-psqlrc', '--csv', '--tuples-only', '--harder'),
         extra_environment=object,
     ).and_return('foo,test,\nbar,test,"stuff and such"')
 
@@ -433,7 +434,7 @@ def test_restore_database_dump_runs_pg_restore():
         extra_environment={'PGSSLMODE': 'disable'},
     ).once()
     flexmock(module).should_receive('execute_command').with_args(
-        ('psql', '--no-password', '--quiet', '--dbname', 'foo', '--command', 'ANALYZE'),
+        ('psql', '--no-password', '--no-psqlrc', '--quiet', '--dbname', 'foo', '--command', 'ANALYZE'),
         extra_environment={'PGSSLMODE': 'disable'},
     ).once()
 
@@ -489,6 +490,7 @@ def test_restore_database_dump_runs_pg_restore_with_hostname_and_port():
         (
             'psql',
             '--no-password',
+            '--no-psqlrc',
             '--quiet',
             '--host',
             'database.example.org',
@@ -539,6 +541,7 @@ def test_restore_database_dump_runs_pg_restore_with_username_and_password():
         (
             'psql',
             '--no-password',
+            '--no-psqlrc',
             '--quiet',
             '--username',
             'postgres',
@@ -589,6 +592,7 @@ def test_restore_database_dump_runs_pg_restore_with_options():
         (
             'psql',
             '--no-password',
+            '--no-psqlrc',
             '--quiet',
             '--dbname',
             'foo',
@@ -612,14 +616,14 @@ def test_restore_database_dump_runs_psql_for_all_database_dump():
     flexmock(module).should_receive('make_dump_path')
     flexmock(module.dump).should_receive('make_database_dump_filename')
     flexmock(module).should_receive('execute_command_with_processes').with_args(
-        ('psql', '--no-password'),
+        ('psql', '--no-password', '--no-psqlrc', '--set', 'ON_ERROR_STOP=on'),
         processes=[extract_process],
         output_log_level=logging.DEBUG,
         input_file=extract_process.stdout,
         extra_environment={'PGSSLMODE': 'disable'},
     ).once()
     flexmock(module).should_receive('execute_command').with_args(
-        ('psql', '--no-password', '--quiet', '--command', 'ANALYZE'),
+        ('psql', '--no-password', '--no-psqlrc', '--quiet', '--command', 'ANALYZE'),
         extra_environment={'PGSSLMODE': 'disable'},
     ).once()
 
@@ -658,7 +662,7 @@ def test_restore_database_dump_runs_non_default_pg_restore_and_psql():
         extra_environment={'PGSSLMODE': 'disable'},
     ).once()
     flexmock(module).should_receive('execute_command').with_args(
-        ('special_psql', '--no-password', '--quiet', '--dbname', 'foo', '--command', 'ANALYZE'),
+        ('special_psql', '--no-password', '--no-psqlrc', '--quiet', '--dbname', 'foo', '--command', 'ANALYZE'),
         extra_environment={'PGSSLMODE': 'disable'},
     ).once()
 
@@ -703,7 +707,7 @@ def test_restore_database_dump_without_extract_process_restores_from_disk():
         extra_environment={'PGSSLMODE': 'disable'},
     ).once()
     flexmock(module).should_receive('execute_command').with_args(
-        ('psql', '--no-password', '--quiet', '--dbname', 'foo', '--command', 'ANALYZE'),
+        ('psql', '--no-password', '--no-psqlrc', '--quiet', '--dbname', 'foo', '--command', 'ANALYZE'),
         extra_environment={'PGSSLMODE': 'disable'},
     ).once()