From f273e82d74d31f4706c3fcc721a5a97a185a539b Mon Sep 17 00:00:00 2001 From: Divyansh Singh Date: Sat, 15 Apr 2023 02:57:51 +0530 Subject: [PATCH] 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 ebd7b8eb..b89ca761 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 d296fa1c..a982209c 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 8255c87d..667a23c3 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 44e427f1..77b830b7 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 349c04be..70cff920 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 + )