def _local_restore(db_name_or_dump_key, *, conn_db, temp_db, curr_db, prev_db):
"""
Performs a restore using a local database
"""
if db_name_or_dump_key == ':current':
local_restore_db = curr_db
elif db_name_or_dump_key == ':previous':
local_restore_db = prev_db
else:
local_restore_db = database.make_config(db_name_or_dump_key[1:])
if not _db_exists(local_restore_db):
raise RuntimeError(
f'local database {local_restore_db["NAME"]} does not exist.'
' Use "pgclone ls --local" to see local database keys.')
# Perform the local restore process. It is completely valid to
# try to restore the temp_db, so do nothing if this database
# is provided by the user
if local_restore_db != temp_db: # pragma: no branch
_drop_db(temp_db)
create_temp_sql = f'''
CREATE DATABASE {temp_db["NAME"]}
TEMPLATE {local_restore_db["NAME"]}
'''
command.run_psql(create_temp_sql, db=conn_db)
_set_search_path(temp_db, conn_db)
return db_name_or_dump_key
def _remote_restore(db_name_or_dump_key, *, temp_db, conn_db):
# We are restoring from a remote dump. If the dump key is not valid,
# assume it is the latest dump of a database name and get the latest
# dump key
dump_key = db_name_or_dump_key
if not is_valid_dump_key(dump_key):
dump_key = get_latest_dump_key(db_name_or_dump_key)
if not dump_key:
raise RuntimeError('Could not find a dump for database name'
f' "{db_name_or_dump_key}"')
storage_location = settings.get_storage_location()
file_path = os.path.join(storage_location, dump_key)
print_msg(f'Creating the temporary restore db')
_drop_db(temp_db)
create_temp_sql = f'CREATE DATABASE "{temp_db["NAME"]}"'
command.run_psql(create_temp_sql, db=conn_db)
_set_search_path(temp_db, conn_db)
print_msg(f'Running pg_restore on "{dump_key}"')
pg_restore_cmd = ('pg_restore --verbose --no-acl --no-owner '
f'-d \'{database.get_url(temp_db)}\'')
# When restoring, we need to ignore errors because there are certain
# errors we cannot get around when pg restoring some DBs (like Aurora).
# In the future, we may parse the output of the pg_restore command to see
# if an unexpected error happened.
if storage_location.startswith('s3://'): # pragma: no cover
command.run_shell(f'aws s3 cp {file_path} - | {pg_restore_cmd}',
ignore_errors=True)
else:
command.run_shell(f'{pg_restore_cmd} {file_path}', ignore_errors=True)
return dump_key
def _kill_connections_to_database(db):
conn_db = database.make_config('postgres')
kill_connections_sql = f'''
SELECT pg_terminate_backend(pg_stat_activity.pid)
FROM pg_stat_activity
WHERE pg_stat_activity.datname = '{db["NAME"]}'
AND pid <> pg_backend_pid()
'''
command.run_psql(kill_connections_sql, db=conn_db)
def _set_search_path(db, conn_db):
"""
pg_restore does not restore the original search_path variable of the
database, which can cause issues with extensions. This function obtains
the original search path and sets it as the search path of the restored
database
https://www.postgresql.org/message-id/CAKFQuwbWpd1DXA_YY3zXRKGKKLL1ZTi9MHDQ15zJ8ufVBwVJEA%40mail.gmail.com
"""
with connection.cursor() as cursor:
cursor.execute('SHOW search_path;')
search_path = cursor.fetchone()[0]
set_search_path_sql = (f'ALTER DATABASE "{db["NAME"]}"'
f' SET search_path to {search_path}')
command.run_psql(set_search_path_sql, db=conn_db)
def _drop_db(db):
conn_db = database.make_config('postgres')
_kill_connections_to_database(db)
drop_sql = f'DROP DATABASE IF EXISTS "{db["NAME"]}"'
command.run_psql(drop_sql, db=conn_db)
def restore(
db_name_or_dump_key,
pre_swap_hooks=None,
restore_config_name=None,
reversible=False,
):
"""
Restores a database dump
Args:
db_name_or_dump_key (str): When a database name
is provided, finds the most recent dump for that
database and restores it. Restores a specific dump
when a dump key is provided.
pre_swap_hooks (List[str], default=None): The list of pre_swap hooks to
run before swapping the temp restore database with the
main database. The strings are management command names.
restore_config_name (str, default=None): The name of the restore
config to use when running the restore. Configured
in settings.PGCLONE_RESTORE_CONFIGS
Returns:
str: The dump key that was restored.
"""
settings.validate()
if not settings.get_allow_restore(): # pragma: no cover
raise RuntimeError('Restore not allowed')
if restore_config_name and pre_swap_hooks: # pragma: no cover
raise ValueError(
'Cannot pass in pre_swap_hooks when using a restore config')
elif pre_swap_hooks:
restore_config_name = '_custom'
restore_config = _make_restore_config(pre_swap_hooks=pre_swap_hooks)
else:
restore_config_name = restore_config_name or 'default'
restore_config = _get_restore_config(restore_config_name)
# Restore works in the following steps with the following databases:
# 1. Create the temp_db database to perform the restore without
# affecting the default_db
# 2. Call pg_restore on temp_db
# 3. Apply any pre_swap_hooks to temp_db
# 4. Terminate all connections on default_db so that we
# can swap in the temp_db
# 4. Rename default_db to prev_db
# 5. Rename temp_db to default_db
# 6. Delete prev_db
#
# Database variable names below reflect this process. We use the
# "postgres" database (i.e. conn_db) as the connection database
# when using psql for most of these commands
default_db = database.get_default_config()
temp_db = database.make_config(default_db['NAME'] + '__temp')
prev_db = database.make_config(default_db['NAME'] + '__prev')
curr_db = database.make_config(default_db['NAME'] + '__curr')
conn_db = database.make_config('postgres')
local_restore_db = None
if db_name_or_dump_key.startswith(':'):
dump_key = _local_restore(
db_name_or_dump_key,
conn_db=conn_db,
temp_db=temp_db,
curr_db=curr_db,
prev_db=prev_db,
)
else:
dump_key = _remote_restore(db_name_or_dump_key,
temp_db=temp_db,
conn_db=conn_db)
# When in reversible mode, make a special __curr db snapshot.
# Avoid this if we are restoring the current db
if reversible and local_restore_db != curr_db:
print_msg('Creating snapshot for reversible restore')
_drop_db(curr_db)
create_current_db_sql = f'''
CREATE DATABASE "{curr_db['NAME']}"
WITH TEMPLATE
"{temp_db['NAME']}"
'''
command.run_psql(create_current_db_sql, db=conn_db)
# pre-swap hook step
with pgconnection.route(temp_db):
for management_command_name in restore_config['pre_swap_hooks']:
print_msg(
f'Running "manage.py {management_command_name}" pre_swap hook')
command.run_management(management_command_name)
# swap step
print_msg('Swapping the restored copy with the primary database')
_drop_db(prev_db)
_kill_connections_to_database(default_db)
alter_db_sql = f'''
ALTER DATABASE "{default_db['NAME']}" RENAME TO
"{prev_db['NAME']}"
'''
# There's a scenario where the default DB may not exist before running
# this, so just ignore errors on this command
command.run_psql(alter_db_sql, db=conn_db, ignore_errors=True)
rename_sql = f'''
ALTER DATABASE "{temp_db["NAME"]}"
RENAME TO "{default_db["NAME"]}"
'''
command.run_psql(rename_sql, db=conn_db)
if not reversible:
print_msg('Cleaning old pgclone resources')
_drop_db(curr_db)
_drop_db(prev_db)
print_msg(f'Successfully restored dump "{dump_key}"')
return dump_key
