Restore Database from AWS S3

1
import io
2
import logging
3
import getpass
4

5
import singlestoredb as s2
6
from IPython.display import display, HTML

1
aws_key_id = None
2
aws_secret_key = None
3
aws_region = 'us-east-1'
4
aws_session_token = ''
5
target_db_name = None
6
backup_id = None

1
def show_warn(warn_msg):
2
    """
3
    Display a warning message in a formatted HTML alert box.
4

5
    Parameters
6
    ----------
7
    warn_msg : str
8
        The warning message to display.
9
    """
10
    display(HTML(f'''<div class="alert alert-block alert-warning">
11
    <b class="fa fa-solid fa-exclamation-circle"></b>
12
    <div>
13
        <p><b>Action Required</b></p>
14
        <p>{warn_msg}</p>
15
    </div>
16
</div>'''))
17

18

19
def show_error(error_msg):
20
    """
21
    Display an error message in a formatted HTML alert box.
22

23
    Parameters
24
    ----------
25
    error_msg : str
26
        The error message to display.
27
    """
28
    display(HTML(f'''<div class="alert alert-block alert-danger">
29
    <b class="fa fa-solid fa-exclamation-triangle"></b>
30
    <div>
31
        <p><b>Error</b></p>
32
        <p>{error_msg}</p>
33
    </div>
34
</div>'''))
35

36

37
def show_success(success_msg):
38
    """
39
    Display a success message in a formatted HTML alert box.
40

41
    Parameters
42
    ----------
43
    success_msg : str
44
        The success message to display.
45
    """
46
    display(HTML(f'''<div class="alert alert-block alert-success">
47
    <b class="fa fa-solid fa-check-circle"></b>
48
    <div>
49
        <p><b>Success</b></p>
50
        <p>{success_msg}</p>
51
    </div>
52
</div>'''))

1
def enable_debug_logs(enabled):
2
    if enabled:
3
        logging.getLogger().setLevel(logging.DEBUG)
4
    else:
5
        logging.getLogger().setLevel(logging.ERROR)

1
def get_sql_statement(p_database_name, p_s3_backup_path):
2

3
    global aws_key_id, aws_secret_key, aws_region, target_db_name, aws_session_token
4
    aws_key_id = (input('Enter AWS_API_KEY_ID:') if aws_key_id == None else aws_key_id)
5
    aws_secret_key = (getpass.getpass('Enter AWS_API_SECRET:') if aws_secret_key == None else aws_secret_key)
6
    aws_session_token = (input('Enter AWS_SESSION_TOKEN:') if aws_session_token == None else aws_session_token)
7
    aws_region = (input('Enter AWS_REGION:') if aws_region == None else aws_region)
8
    target_db_name = (input('Enter Target Database Name:') if target_db_name == None else target_db_name)
9

10

11

12
    data = io.StringIO()
13
    data.write('RESTORE DATABASE ' + p_database_name + ' ')
14

15
    if target_db_name != '':
16
        data.write(' AS ' + target_db_name + ' ')
17

18
    data.write(' FROM S3 "' + p_s3_backup_path + '" ')
19
    data.write(' CONFIG \' {"region":"' + aws_region + '"} \'')
20
    data.write(' CREDENTIALS \'{"aws_access_key_id":"' + aws_key_id
21
               + '","aws_secret_access_key":"' + aws_secret_key + '"')
22
    if aws_session_token != '':
23
        data.write(',"aws_session_token":"' + aws_session_token +'" ')
24
    data.write('}\' ')
25
    logging.debug(f'statement: {data.getvalue()}')
26
    return data.getvalue()

1
# Disable debug logs
2
enable_debug_logs(False)
3

4
# Start of script execution
5
print('Script execution began')
6

7
# Check if the connection URL ends with '/'
8
if connection_url.endswith('/'):
9
    # Show a warning message if the database is not selected
10
    show_warn('Database not selected. Please select from dropdown in top of web page')
11
else:
12
    try:
13
        # Get the backup ID from the secret store
14
        backup_id = (input('Enter Backup ID:') if backup_id == None else backup_id)
15
        backup_db_name = None
16
        # Establish a connection to the database
17
        conn = s2.connect(results_type='dict')
18
        with conn.cursor() as cursor:
19
            # Fetch the database name and backup path from the backup history
20
            query = f'SELECT DATABASE_NAME, BACKUP_PATH FROM information_schema.mv_backup_history WHERE BACKUP_ID = {backup_id} and STATUS = \'Success\''
21
            logging.debug(f'query: {query}')
22
            cursor.execute(query)
23
            record = cursor.fetchone()
24
            if record is None:
25
                # Log an error and show an error message if no record is found for the specified backup ID
26
                logging.error('No record found for the specified backup ID')
27
                show_error('No record found for the specified backup ID')
28
            else:
29
                # Get the database name and backup path from the fetched record
30
                backup_db_name = record['DATABASE_NAME']
31
                backup_path = record['BACKUP_PATH']
32

33
                # Execute the SQL statement to restore the database
34
                cursor.execute(get_sql_statement(p_database_name=backup_db_name, p_s3_backup_path=backup_path))
35
                results = cursor.fetchall()
36

37
                # Check if the restore was successful
38
                if results is None:
39
                    logging.error('Restore execution failed')
40
                    show_error('Restore Failed')
41
                else:
42
                    logging.info("Restore completed")
43
                    verify_query = f"select schema_name from information_schema.SCHEMATA where schema_name = {target_db_name if target_db_name != '' else backup_db_name}"
44
                    show_success(f'Restore completed.\n {verify_query}')
45

46
    except s2.exceptions.OperationalError as ope:
47
        # Handle specific operational errors
48
        if 'NoSuchBucket' in ope.errmsg:
49
            logging.error('Provided S3 Bucket does not exist. Please check')
50
            show_error('Provided S3 Bucket does not exist. Please check')
51
        elif 'Access denied' in ope.errmsg:
52
            logging.error('Failed to restore due to grants missing or firewall settings. Please check')
53
            show_error('Failed to restore due to grants missing or firewall settings. Please check')
54
        else:
55
            logging.error(f'Failed. error msg: {ope.errmsg}')
56
            show_error('Restore Failed. error msg: {ope.errmsg}')
57
    except s2.Error as e:
58
        # Handle any other errors
59
        logging.error(f'Encountered exception {e}')
60
        show_error(f'Restore Failed. {str(e)}')
61

62
# End of script execution
63
print('\n\nScript execution completed')

1
%sql select schema_name from information_schema.SCHEMATA;