Override PostgreSQL dump/restore commands via configuration options (#311).

Merge pull request #49 from jpaniagualaconich/specify-pg-dump-restore-commands

borgmatic/config/schema.yaml

@@ -764,6 +764,32 @@ properties:
                             description: |
                                 Path to a certificate revocation list.
                             example: "/root/.postgresql/root.crl"
+                        pg_dump_command:
+                            type: string
+                            description: |
+                                Command to use instead of "pg_dump" or
+                                "pg_dumpall". This can be used to run a specific
+                                pg_dump version (e.g., one inside a running
+                                docker container). Defaults to "pg_dump" for
+                                single database dump or "pg_dumpall" to dump
+                                all databases.
+                            example: docker exec my_pg_container pg_dump
+                        pg_restore_command:
+                            type: string
+                            description: |
+                                Command to use instead of "pg_restore". This
+                                can be used to run a specific pg_restore
+                                version (e.g., one inside a running docker
+                                container). Defaults to "pg_restore".
+                            example: docker exec my_pg_container pg_restore
+                        psql_command:
+                            type: string
+                            description: |
+                                Command to use instead of "psql". This can be
+                                used to run a specific psql version (e.g.,
+                                one inside a running docker container).
+                                Defaults to "psql".
+                            example: docker exec my_pg_container psql
                         options:
                             type: string
                             description: |

borgmatic/hooks/postgresql.py

@@ -56,9 +56,11 @@ def dump_databases(databases, log_prefix, location_config, dry_run):
         )
         all_databases = bool(name == 'all')
         dump_format = database.get('format', 'custom')
+        default_dump_command = 'pg_dumpall' if all_databases else 'pg_dump'
+        dump_command = database.get('pg_dump_command') or default_dump_command
         command = (
             (
-                'pg_dumpall' if all_databases else 'pg_dump',
+                dump_command,
                 '--no-password',
                 '--clean',
                 '--if-exists',
@@ -140,16 +142,18 @@ 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'
     analyze_command = (
-        ('psql', '--no-password', '--quiet')
+        (psql_command, '--no-password', '--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 ())
         + (('--dbname', database['name']) if not all_databases else ())
         + ('--command', 'ANALYZE')
     )
+    pg_restore_command = database.get('pg_restore_command') or 'pg_restore'
     restore_command = (
-        ('psql' if all_databases else 'pg_restore', '--no-password')
+        (psql_command if all_databases else pg_restore_command, '--no-password')
         + (
             ('--if-exists', '--exit-on-error', '--clean', '--dbname', database['name'])
             if not all_databases

tests/unit/hooks/test_postgresql.py

@@ -223,6 +223,36 @@ def test_dump_databases_runs_pg_dumpall_for_all_databases():
     assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == [process]
 
 
+def test_dump_databases_runs_non_default_pg_dump():
+    databases = [{'name': 'foo', 'pg_dump_command': 'special_pg_dump'}]
+    process = flexmock()
+    flexmock(module).should_receive('make_dump_path').and_return('')
+    flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
+        'databases/localhost/foo'
+    )
+    flexmock(module.dump).should_receive('create_named_pipe_for_dump')
+    flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
+
+    flexmock(module).should_receive('execute_command').with_args(
+        (
+            'special_pg_dump',
+            '--no-password',
+            '--clean',
+            '--if-exists',
+            '--format',
+            'custom',
+            'foo',
+            '>',
+            'databases/localhost/foo',
+        ),
+        shell=True,
+        extra_environment={'PGSSLMODE': 'disable'},
+        run_to_completion=False,
+    ).and_return(process).once()
+
+    assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == [process]
+
+
 def test_restore_database_dump_runs_pg_restore():
     database_config = [{'name': 'foo'}]
     extract_process = flexmock(stdout=flexmock())
@@ -388,6 +418,40 @@ def test_restore_database_dump_runs_psql_for_all_database_dump():
     )
 
 
+def test_restore_database_dump_runs_non_default_pg_restore_and_psql():
+    database_config = [
+        {'name': 'foo', 'pg_restore_command': 'special_pg_restore', 'psql_command': 'special_psql'}
+    ]
+    extract_process = flexmock(stdout=flexmock())
+
+    flexmock(module).should_receive('make_dump_path')
+    flexmock(module.dump).should_receive('make_database_dump_filename')
+    flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
+    flexmock(module).should_receive('execute_command_with_processes').with_args(
+        (
+            'special_pg_restore',
+            '--no-password',
+            '--if-exists',
+            '--exit-on-error',
+            '--clean',
+            '--dbname',
+            'foo',
+        ),
+        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(
+        ('special_psql', '--no-password', '--quiet', '--dbname', 'foo', '--command', 'ANALYZE'),
+        extra_environment={'PGSSLMODE': 'disable'},
+    ).once()
+
+    module.restore_database_dump(
+        database_config, 'test.yaml', {}, dry_run=False, extract_process=extract_process
+    )
+
+
 def test_restore_database_dump_with_dry_run_skips_restore():
     database_config = [{'name': 'foo'}]