feat: restore specific schemas (#375).

Merge pull request #67 from diivi/feat/restore-specific-schemas

borgmatic/actions/restore.py

@@ -313,7 +313,7 @@ def run_restore(
                 remote_path,
                 archive_name,
                 found_hook_name or hook_name,
-                found_database,
+                dict(found_database, **{'schemas': restore_arguments.schemas}),
             )
 
     # For any database that weren't found via exact matches in the hooks configuration, try to
@@ -342,7 +342,7 @@ def run_restore(
                 remote_path,
                 archive_name,
                 found_hook_name or hook_name,
-                database,
+                dict(database, **{'schemas': restore_arguments.schemas}),
             )
 
     borgmatic.hooks.dispatch.call_hooks_even_if_unconfigured(

borgmatic/commands/arguments.py

@@ -629,6 +629,13 @@ def make_parsers():
         dest='databases',
         help="Names of databases to restore from archive, defaults to all databases. Note that any databases to restore must be defined in borgmatic's configuration",
     )
+    restore_group.add_argument(
+        '--schema',
+        metavar='NAME',
+        nargs='+',
+        dest='schemas',
+        help='Names of schemas to restore from the database, defaults to all schemas. Schemas are only supported for PostgreSQL and MongoDB databases',
+    )
     restore_group.add_argument(
         '-h', '--help', action='help', help='Show this help message and exit'
     )

borgmatic/hooks/mongodb.py

@@ -161,4 +161,7 @@ def build_restore_command(extract_process, database, dump_filename):
         command.extend(('--authenticationDatabase', database['authentication_database']))
     if 'restore_options' in database:
         command.extend(database['restore_options'].split(' '))
+    if database['schemas']:
+        for schema in database['schemas']:
+            command.extend(('--nsInclude', schema))
     return command

borgmatic/hooks/postgresql.py

@@ -1,4 +1,5 @@
 import csv
+import itertools
 import logging
 import os
 
@@ -225,7 +226,13 @@ def restore_database_dump(database_config, log_prefix, location_config, dry_run,
         + (('--username', database['username']) if 'username' in database else ())
         + (tuple(database['restore_options'].split(' ')) if 'restore_options' in database else ())
         + (() if extract_process else (dump_filename,))
+        + tuple(
+            itertools.chain.from_iterable(('--schema', schema) for schema in database['schemas'])
+            if database['schemas']
+            else ()
         )
+    )
+
     extra_environment = make_extra_environment(database)
 
     logger.debug(f"{log_prefix}: Restoring PostgreSQL database {database['name']}{dry_run_label}")

tests/unit/actions/test_restore.py

@@ -233,7 +233,7 @@ def test_run_restore_restores_each_database():
         remote_path=object,
         archive_name=object,
         hook_name='postgresql_databases',
-        database={'name': 'foo'},
+        database={'name': 'foo', 'schemas': None},
     ).once()
     flexmock(module).should_receive('restore_single_database').with_args(
         repository=object,
@@ -246,7 +246,7 @@ def test_run_restore_restores_each_database():
         remote_path=object,
         archive_name=object,
         hook_name='postgresql_databases',
-        database={'name': 'bar'},
+        database={'name': 'bar', 'schemas': None},
     ).once()
     flexmock(module).should_receive('ensure_databases_found')
 
@@ -256,7 +256,9 @@ def test_run_restore_restores_each_database():
         storage=flexmock(),
         hooks=flexmock(),
         local_borg_version=flexmock(),
-        restore_arguments=flexmock(repository='repo', archive='archive', databases=flexmock()),
+        restore_arguments=flexmock(
+            repository='repo', archive='archive', databases=flexmock(), schemas=None
+        ),
         global_arguments=flexmock(dry_run=False),
         local_path=flexmock(),
         remote_path=flexmock(),
@@ -327,7 +329,7 @@ def test_run_restore_restores_database_configured_with_all_name():
         remote_path=object,
         archive_name=object,
         hook_name='postgresql_databases',
-        database={'name': 'foo'},
+        database={'name': 'foo', 'schemas': None},
     ).once()
     flexmock(module).should_receive('restore_single_database').with_args(
         repository=object,
@@ -340,7 +342,7 @@ def test_run_restore_restores_database_configured_with_all_name():
         remote_path=object,
         archive_name=object,
         hook_name='postgresql_databases',
-        database={'name': 'bar'},
+        database={'name': 'bar', 'schemas': None},
     ).once()
     flexmock(module).should_receive('ensure_databases_found')
 
@@ -350,7 +352,9 @@ def test_run_restore_restores_database_configured_with_all_name():
         storage=flexmock(),
         hooks=flexmock(),
         local_borg_version=flexmock(),
-        restore_arguments=flexmock(repository='repo', archive='archive', databases=flexmock()),
+        restore_arguments=flexmock(
+            repository='repo', archive='archive', databases=flexmock(), schemas=None
+        ),
         global_arguments=flexmock(dry_run=False),
         local_path=flexmock(),
         remote_path=flexmock(),
@@ -399,7 +403,7 @@ def test_run_restore_skips_missing_database():
         remote_path=object,
         archive_name=object,
         hook_name='postgresql_databases',
-        database={'name': 'foo'},
+        database={'name': 'foo', 'schemas': None},
     ).once()
     flexmock(module).should_receive('restore_single_database').with_args(
         repository=object,
@@ -412,7 +416,7 @@ def test_run_restore_skips_missing_database():
         remote_path=object,
         archive_name=object,
         hook_name='postgresql_databases',
-        database={'name': 'bar'},
+        database={'name': 'bar', 'schemas': None},
     ).never()
     flexmock(module).should_receive('ensure_databases_found')
 
@@ -422,7 +426,9 @@ def test_run_restore_skips_missing_database():
         storage=flexmock(),
         hooks=flexmock(),
         local_borg_version=flexmock(),
-        restore_arguments=flexmock(repository='repo', archive='archive', databases=flexmock()),
+        restore_arguments=flexmock(
+            repository='repo', archive='archive', databases=flexmock(), schemas=None
+        ),
         global_arguments=flexmock(dry_run=False),
         local_path=flexmock(),
         remote_path=flexmock(),
@@ -465,7 +471,7 @@ def test_run_restore_restores_databases_from_different_hooks():
         remote_path=object,
         archive_name=object,
         hook_name='postgresql_databases',
-        database={'name': 'foo'},
+        database={'name': 'foo', 'schemas': None},
     ).once()
     flexmock(module).should_receive('restore_single_database').with_args(
         repository=object,
@@ -478,7 +484,7 @@ def test_run_restore_restores_databases_from_different_hooks():
         remote_path=object,
         archive_name=object,
         hook_name='mysql_databases',
-        database={'name': 'bar'},
+        database={'name': 'bar', 'schemas': None},
     ).once()
     flexmock(module).should_receive('ensure_databases_found')
 
@@ -488,7 +494,9 @@ def test_run_restore_restores_databases_from_different_hooks():
         storage=flexmock(),
         hooks=flexmock(),
         local_borg_version=flexmock(),
-        restore_arguments=flexmock(repository='repo', archive='archive', databases=flexmock()),
+        restore_arguments=flexmock(
+            repository='repo', archive='archive', databases=flexmock(), schemas=None
+        ),
         global_arguments=flexmock(dry_run=False),
         local_path=flexmock(),
         remote_path=flexmock(),

tests/unit/hooks/test_mongodb.py

@@ -157,7 +157,7 @@ def test_dump_databases_runs_mongodumpall_for_all_databases():
 
 
 def test_restore_database_dump_runs_mongorestore():
-    database_config = [{'name': 'foo'}]
+    database_config = [{'name': 'foo', 'schemas': None}]
     extract_process = flexmock(stdout=flexmock())
 
     flexmock(module).should_receive('make_dump_path')
@@ -189,7 +189,9 @@ def test_restore_database_dump_errors_on_multiple_database_config():
 
 
 def test_restore_database_dump_runs_mongorestore_with_hostname_and_port():
-    database_config = [{'name': 'foo', 'hostname': 'database.example.org', 'port': 5433}]
+    database_config = [
+        {'name': 'foo', 'hostname': 'database.example.org', 'port': 5433, 'schemas': None}
+    ]
     extract_process = flexmock(stdout=flexmock())
 
     flexmock(module).should_receive('make_dump_path')
@@ -223,6 +225,7 @@ def test_restore_database_dump_runs_mongorestore_with_username_and_password():
             'username': 'mongo',
             'password': 'trustsome1',
             'authentication_database': 'admin',
+            'schemas': None,
         }
     ]
     extract_process = flexmock(stdout=flexmock())
@@ -254,7 +257,7 @@ def test_restore_database_dump_runs_mongorestore_with_username_and_password():
 
 
 def test_restore_database_dump_runs_mongorestore_with_options():
-    database_config = [{'name': 'foo', 'restore_options': '--harder'}]
+    database_config = [{'name': 'foo', 'restore_options': '--harder', 'schemas': None}]
     extract_process = flexmock(stdout=flexmock())
 
     flexmock(module).should_receive('make_dump_path')
@@ -271,8 +274,36 @@ def test_restore_database_dump_runs_mongorestore_with_options():
     )
 
 
+def test_restore_databases_dump_runs_mongorestore_with_schemas():
+    database_config = [{'name': 'foo', 'schemas': ['bar', 'baz']}]
+    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('execute_command_with_processes').with_args(
+        [
+            'mongorestore',
+            '--archive',
+            '--drop',
+            '--db',
+            'foo',
+            '--nsInclude',
+            'bar',
+            '--nsInclude',
+            'baz',
+        ],
+        processes=[extract_process],
+        output_log_level=logging.DEBUG,
+        input_file=extract_process.stdout,
+    ).once()
+
+    module.restore_database_dump(
+        database_config, 'test.yaml', {}, dry_run=False, extract_process=extract_process
+    )
+
+
 def test_restore_database_dump_runs_psql_for_all_database_dump():
-    database_config = [{'name': 'all'}]
+    database_config = [{'name': 'all', 'schemas': None}]
     extract_process = flexmock(stdout=flexmock())
 
     flexmock(module).should_receive('make_dump_path')
@@ -290,7 +321,7 @@ def test_restore_database_dump_runs_psql_for_all_database_dump():
 
 
 def test_restore_database_dump_with_dry_run_skips_restore():
-    database_config = [{'name': 'foo'}]
+    database_config = [{'name': 'foo', 'schemas': None}]
 
     flexmock(module).should_receive('make_dump_path')
     flexmock(module.dump).should_receive('make_database_dump_filename')
@@ -302,7 +333,7 @@ def test_restore_database_dump_with_dry_run_skips_restore():
 
 
 def test_restore_database_dump_without_extract_process_restores_from_disk():
-    database_config = [{'name': 'foo', 'format': 'directory'}]
+    database_config = [{'name': 'foo', 'format': 'directory', 'schemas': None}]
 
     flexmock(module).should_receive('make_dump_path')
     flexmock(module.dump).should_receive('make_database_dump_filename').and_return('/dump/path')

tests/unit/hooks/test_postgresql.py

@@ -411,7 +411,7 @@ def test_dump_databases_runs_non_default_pg_dump():
 
 
 def test_restore_database_dump_runs_pg_restore():
-    database_config = [{'name': 'foo'}]
+    database_config = [{'name': 'foo', 'schemas': None}]
     extract_process = flexmock(stdout=flexmock())
 
     flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
@@ -458,7 +458,9 @@ def test_restore_database_dump_errors_on_multiple_database_config():
 
 
 def test_restore_database_dump_runs_pg_restore_with_hostname_and_port():
-    database_config = [{'name': 'foo', 'hostname': 'database.example.org', 'port': 5433}]
+    database_config = [
+        {'name': 'foo', 'hostname': 'database.example.org', 'port': 5433, 'schemas': None}
+    ]
     extract_process = flexmock(stdout=flexmock())
 
     flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
@@ -506,7 +508,9 @@ def test_restore_database_dump_runs_pg_restore_with_hostname_and_port():
 
 
 def test_restore_database_dump_runs_pg_restore_with_username_and_password():
-    database_config = [{'name': 'foo', 'username': 'postgres', 'password': 'trustsome1'}]
+    database_config = [
+        {'name': 'foo', 'username': 'postgres', 'password': 'trustsome1', 'schemas': None}
+    ]
     extract_process = flexmock(stdout=flexmock())
 
     flexmock(module).should_receive('make_extra_environment').and_return(
@@ -553,7 +557,12 @@ def test_restore_database_dump_runs_pg_restore_with_username_and_password():
 
 def test_restore_database_dump_runs_pg_restore_with_options():
     database_config = [
-        {'name': 'foo', 'restore_options': '--harder', 'analyze_options': '--smarter'}
+        {
+            'name': 'foo',
+            'restore_options': '--harder',
+            'analyze_options': '--smarter',
+            'schemas': None,
+        }
     ]
     extract_process = flexmock(stdout=flexmock())
 
@@ -596,7 +605,7 @@ def test_restore_database_dump_runs_pg_restore_with_options():
 
 
 def test_restore_database_dump_runs_psql_for_all_database_dump():
-    database_config = [{'name': 'all'}]
+    database_config = [{'name': 'all', 'schemas': None}]
     extract_process = flexmock(stdout=flexmock())
 
     flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
@@ -621,7 +630,12 @@ 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'}
+        {
+            'name': 'foo',
+            'pg_restore_command': 'special_pg_restore',
+            'psql_command': 'special_psql',
+            'schemas': None,
+        }
     ]
     extract_process = flexmock(stdout=flexmock())
 
@@ -654,7 +668,7 @@ def test_restore_database_dump_runs_non_default_pg_restore_and_psql():
 
 
 def test_restore_database_dump_with_dry_run_skips_restore():
-    database_config = [{'name': 'foo'}]
+    database_config = [{'name': 'foo', 'schemas': None}]
 
     flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
     flexmock(module).should_receive('make_dump_path')
@@ -667,7 +681,7 @@ def test_restore_database_dump_with_dry_run_skips_restore():
 
 
 def test_restore_database_dump_without_extract_process_restores_from_disk():
-    database_config = [{'name': 'foo'}]
+    database_config = [{'name': 'foo', 'schemas': None}]
 
     flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
     flexmock(module).should_receive('make_dump_path')
@@ -696,3 +710,39 @@ def test_restore_database_dump_without_extract_process_restores_from_disk():
     module.restore_database_dump(
         database_config, 'test.yaml', {}, dry_run=False, extract_process=None
     )
+
+
+def test_restore_database_dump_with_schemas_restores_schemas():
+    database_config = [{'name': 'foo', 'schemas': ['bar', 'baz']}]
+
+    flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
+    flexmock(module).should_receive('make_dump_path')
+    flexmock(module.dump).should_receive('make_database_dump_filename').and_return('/dump/path')
+    flexmock(module).should_receive('execute_command_with_processes').with_args(
+        (
+            'pg_restore',
+            '--no-password',
+            '--if-exists',
+            '--exit-on-error',
+            '--clean',
+            '--dbname',
+            'foo',
+            '/dump/path',
+            '--schema',
+            'bar',
+            '--schema',
+            'baz',
+        ),
+        processes=[],
+        output_log_level=logging.DEBUG,
+        input_file=None,
+        extra_environment={'PGSSLMODE': 'disable'},
+    ).once()
+    flexmock(module).should_receive('execute_command').with_args(
+        ('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=None
+    )