Restore Database from AWS S3

import io
import logging
import getpass

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

aws_key_id = None
aws_secret_key = None
aws_region = 'us-east-1'
aws_session_token = ''
target_db_name = None
backup_id = None

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

    Parameters
    ----------
    warn_msg : str
        The warning message to display.
    """
    display(HTML(f'''<div class="alert alert-block alert-warning">
    <b class="fa fa-solid fa-exclamation-circle"></b>
    <div>
        <p><b>Action Required</b></p>
        <p>{warn_msg}</p>
    </div>
</div>'''))


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

    Parameters
    ----------
    error_msg : str
        The error message to display.
    """
    display(HTML(f'''<div class="alert alert-block alert-danger">
    <b class="fa fa-solid fa-exclamation-triangle"></b>
    <div>
        <p><b>Error</b></p>
        <p>{error_msg}</p>
    </div>
</div>'''))


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

    Parameters
    ----------
    success_msg : str
        The success message to display.
    """
    display(HTML(f'''<div class="alert alert-block alert-success">
    <b class="fa fa-solid fa-check-circle"></b>
    <div>
        <p><b>Success</b></p>
        <p>{success_msg}</p>
    </div>
</div>'''))

def enable_debug_logs(enabled):
    if enabled:
        logging.getLogger().setLevel(logging.DEBUG)
    else:
        logging.getLogger().setLevel(logging.ERROR)

def get_sql_statement(p_database_name, p_s3_backup_path):

    global aws_key_id, aws_secret_key, aws_region, target_db_name, aws_session_token
    aws_key_id = (input('Enter AWS_API_KEY_ID:') if aws_key_id == None else aws_key_id)
    aws_secret_key = (getpass.getpass('Enter AWS_API_SECRET:') if aws_secret_key == None else aws_secret_key)
    aws_session_token = (input('Enter AWS_SESSION_TOKEN:') if aws_session_token == None else aws_session_token)
    aws_region = (input('Enter AWS_REGION:') if aws_region == None else aws_region)
    target_db_name = (input('Enter Target Database Name:') if target_db_name == None else target_db_name)



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

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

    data.write(' FROM S3 "' + p_s3_backup_path + '" ')
    data.write(' CONFIG \' {"region":"' + aws_region + '"} \'')
    data.write(' CREDENTIALS \'{"aws_access_key_id":"' + aws_key_id
               + '","aws_secret_access_key":"' + aws_secret_key + '"')
    if aws_session_token != '':
        data.write(',"aws_session_token":"' + aws_session_token +'" ')
    data.write('}\' ')
    logging.debug(f'statement: {data.getvalue()}')
    return data.getvalue()

# Disable debug logs
enable_debug_logs(False)

# Start of script execution
print('Script execution began')

# Check if the connection URL ends with '/'
if connection_url.endswith('/'):
    # Show a warning message if the database is not selected
    show_warn('Database not selected. Please select from dropdown in top of web page')
else:
    try:
        # Get the backup ID from the secret store
        backup_id = (input('Enter Backup ID:') if backup_id == None else backup_id)
        backup_db_name = None
        # Establish a connection to the database
        conn = s2.connect(results_type='dict')
        with conn.cursor() as cursor:
            # Fetch the database name and backup path from the backup history
            query = f'SELECT DATABASE_NAME, BACKUP_PATH FROM information_schema.mv_backup_history WHERE BACKUP_ID = {backup_id} and STATUS = \'Success\''
            logging.debug(f'query: {query}')
            cursor.execute(query)
            record = cursor.fetchone()
            if record is None:
                # Log an error and show an error message if no record is found for the specified backup ID
                logging.error('No record found for the specified backup ID')
                show_error('No record found for the specified backup ID')
            else:
                # Get the database name and backup path from the fetched record
                backup_db_name = record['DATABASE_NAME']
                backup_path = record['BACKUP_PATH']

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

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

    except s2.exceptions.OperationalError as ope:
        # Handle specific operational errors
        if 'NoSuchBucket' in ope.errmsg:
            logging.error('Provided S3 Bucket does not exist. Please check')
            show_error('Provided S3 Bucket does not exist. Please check')
        elif 'Access denied' in ope.errmsg:
            logging.error('Failed to restore due to grants missing or firewall settings. Please check')
            show_error('Failed to restore due to grants missing or firewall settings. Please check')
        else:
            logging.error(f'Failed. error msg: {ope.errmsg}')
            show_error('Restore Failed. error msg: {ope.errmsg}')
    except s2.Error as e:
        # Handle any other errors
        logging.error(f'Encountered exception {e}')
        show_error(f'Restore Failed. {str(e)}')

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

%sql select schema_name from information_schema.SCHEMATA;