diff --git a/borgmatic/actions/restore.py b/borgmatic/actions/restore.py index bbbb4c7..1e8d175 100644 --- a/borgmatic/actions/restore.py +++ b/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( diff --git a/borgmatic/commands/arguments.py b/borgmatic/commands/arguments.py index 13ee315..b89ca76 100644 --- a/borgmatic/commands/arguments.py +++ b/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' ) diff --git a/borgmatic/hooks/mongodb.py b/borgmatic/hooks/mongodb.py index be5f656..781e5f2 100644 --- a/borgmatic/hooks/mongodb.py +++ b/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 diff --git a/borgmatic/hooks/postgresql.py b/borgmatic/hooks/postgresql.py index d4799f5..a982209 100644 --- a/borgmatic/hooks/postgresql.py +++ b/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}") diff --git a/tests/unit/actions/test_restore.py b/tests/unit/actions/test_restore.py index 8255c87..667a23c 100644 --- a/tests/unit/actions/test_restore.py +++ b/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(), diff --git a/tests/unit/hooks/test_mongodb.py b/tests/unit/hooks/test_mongodb.py index 44e427f..77b830b 100644 --- a/tests/unit/hooks/test_mongodb.py +++ b/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') diff --git a/tests/unit/hooks/test_postgresql.py b/tests/unit/hooks/test_postgresql.py index 349c04b..70cff92 100644 --- a/tests/unit/hooks/test_postgresql.py +++ b/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 + )