feat: restore specific schemas (#375).

Merge pull request #67 from diivi/feat/restore-specific-schemas
This commit is contained in:
Dan Helfman 2023-04-14 16:26:25 -07:00 committed by GitHub
commit 81e167959b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 134 additions and 28 deletions

View File

@ -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(

View File

@ -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'
)

View File

@ -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

View File

@ -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}")

View File

@ -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(),

View File

@ -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')

View File

@ -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
)