From 9bc2322f9a38c6ab704be284563b1928c7673172 Mon Sep 17 00:00:00 2001 From: Divyansh Singh Date: Thu, 6 Apr 2023 02:10:36 +0530 Subject: [PATCH 1/4] feat: restore specific schemas --- borgmatic/actions/restore.py | 5 +++++ borgmatic/commands/arguments.py | 7 +++++++ borgmatic/hooks/postgresql.py | 2 ++ 3 files changed, 14 insertions(+) diff --git a/borgmatic/actions/restore.py b/borgmatic/actions/restore.py index bbbb4c7..fa7b04b 100644 --- a/borgmatic/actions/restore.py +++ b/borgmatic/actions/restore.py @@ -68,12 +68,15 @@ def restore_single_database( archive_name, hook_name, database, + schemas ): # pragma: no cover ''' Given (among other things) an archive name, a database hook name, and a configured database configuration dict, restore that database from the archive. ''' logger.info(f'{repository}: Restoring database {database["name"]}') + if schemas: + database['schemas'] = schemas dump_pattern = borgmatic.hooks.dispatch.call_hooks( 'make_database_dump_pattern', @@ -314,6 +317,7 @@ def run_restore( archive_name, found_hook_name or hook_name, found_database, + schemas = restore_arguments.schemas, ) # For any database that weren't found via exact matches in the hooks configuration, try to @@ -343,6 +347,7 @@ def run_restore( archive_name, found_hook_name or hook_name, 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..7f93910 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." + ) restore_group.add_argument( '-h', '--help', action='help', help='Show this help message and exit' ) diff --git a/borgmatic/hooks/postgresql.py b/borgmatic/hooks/postgresql.py index d4799f5..5d951f6 100644 --- a/borgmatic/hooks/postgresql.py +++ b/borgmatic/hooks/postgresql.py @@ -213,6 +213,7 @@ def restore_database_dump(database_config, log_prefix, location_config, dry_run, + ('--command', 'ANALYZE') ) pg_restore_command = database.get('pg_restore_command') or 'pg_restore' + backup_schemas = ', '.join(database['schemas']) if 'schemas' in database else None restore_command = ( (psql_command if all_databases else pg_restore_command, '--no-password') + ( @@ -223,6 +224,7 @@ def restore_database_dump(database_config, log_prefix, location_config, dry_run, + (('--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 ()) + + (('--schema', backup_schemas) if backup_schemas else ()) + (tuple(database['restore_options'].split(' ')) if 'restore_options' in database else ()) + (() if extract_process else (dump_filename,)) ) From 264cebd2b10cd13c3782d639bba32e07f556c2be Mon Sep 17 00:00:00 2001 From: Divyansh Singh Date: Tue, 11 Apr 2023 23:19:49 +0530 Subject: [PATCH 2/4] complete psql multi schema backup --- borgmatic/actions/restore.py | 9 ++------- borgmatic/commands/arguments.py | 2 +- borgmatic/hooks/postgresql.py | 7 +++++-- 3 files changed, 8 insertions(+), 10 deletions(-) diff --git a/borgmatic/actions/restore.py b/borgmatic/actions/restore.py index fa7b04b..1e8d175 100644 --- a/borgmatic/actions/restore.py +++ b/borgmatic/actions/restore.py @@ -68,15 +68,12 @@ def restore_single_database( archive_name, hook_name, database, - schemas ): # pragma: no cover ''' Given (among other things) an archive name, a database hook name, and a configured database configuration dict, restore that database from the archive. ''' logger.info(f'{repository}: Restoring database {database["name"]}') - if schemas: - database['schemas'] = schemas dump_pattern = borgmatic.hooks.dispatch.call_hooks( 'make_database_dump_pattern', @@ -316,8 +313,7 @@ def run_restore( remote_path, archive_name, found_hook_name or hook_name, - found_database, - schemas = restore_arguments.schemas, + dict(found_database, **{'schemas': restore_arguments.schemas}), ) # For any database that weren't found via exact matches in the hooks configuration, try to @@ -346,8 +342,7 @@ def run_restore( remote_path, archive_name, found_hook_name or hook_name, - database, - schemas = restore_arguments.schemas, + 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 7f93910..ebd7b8e 100644 --- a/borgmatic/commands/arguments.py +++ b/borgmatic/commands/arguments.py @@ -634,7 +634,7 @@ def make_parsers(): metavar='NAME', nargs='+', dest='schemas', - help="Names of schemas to restore from the database, defaults to all 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/postgresql.py b/borgmatic/hooks/postgresql.py index 5d951f6..d296fa1 100644 --- a/borgmatic/hooks/postgresql.py +++ b/borgmatic/hooks/postgresql.py @@ -213,7 +213,6 @@ def restore_database_dump(database_config, log_prefix, location_config, dry_run, + ('--command', 'ANALYZE') ) pg_restore_command = database.get('pg_restore_command') or 'pg_restore' - backup_schemas = ', '.join(database['schemas']) if 'schemas' in database else None restore_command = ( (psql_command if all_databases else pg_restore_command, '--no-password') + ( @@ -224,10 +223,14 @@ def restore_database_dump(database_config, log_prefix, location_config, dry_run, + (('--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 ()) - + (('--schema', backup_schemas) if backup_schemas else ()) + (tuple(database['restore_options'].split(' ')) if 'restore_options' in database else ()) + (() if extract_process else (dump_filename,)) ) + + if database['schemas']: + for schema in database['schemas']: + restore_command += ('--schema', schema) + extra_environment = make_extra_environment(database) logger.debug(f"{log_prefix}: Restoring PostgreSQL database {database['name']}{dry_run_label}") From 2fea429d785ba9902057472c2d3afc2cdb7ffa79 Mon Sep 17 00:00:00 2001 From: Divyansh Singh Date: Wed, 12 Apr 2023 09:34:19 +0530 Subject: [PATCH 3/4] collection restore for mongodb --- borgmatic/hooks/mongodb.py | 3 +++ 1 file changed, 3 insertions(+) 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 From f273e82d74d31f4706c3fcc721a5a97a185a539b Mon Sep 17 00:00:00 2001 From: Divyansh Singh Date: Sat, 15 Apr 2023 02:57:51 +0530 Subject: [PATCH 4/4] add tests --- borgmatic/commands/arguments.py | 2 +- borgmatic/hooks/postgresql.py | 10 +++-- tests/unit/actions/test_restore.py | 32 ++++++++------ tests/unit/hooks/test_mongodb.py | 43 ++++++++++++++++--- tests/unit/hooks/test_postgresql.py | 66 +++++++++++++++++++++++++---- 5 files changed, 122 insertions(+), 31 deletions(-) diff --git a/borgmatic/commands/arguments.py b/borgmatic/commands/arguments.py index ebd7b8e..b89ca76 100644 --- a/borgmatic/commands/arguments.py +++ b/borgmatic/commands/arguments.py @@ -634,7 +634,7 @@ def make_parsers(): 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", + 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/postgresql.py b/borgmatic/hooks/postgresql.py index d296fa1..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,12 +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 () + ) ) - if database['schemas']: - for schema in database['schemas']: - restore_command += ('--schema', schema) - 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 + )