Backup Database to AWS S3

1import io2import logging3import time4import getpass5
6import singlestoredb as s27from IPython.display import display, HTML

1aws_key_id = None2aws_secret_key = None3aws_region = 'us-east-1'4s3_target_path = None5aws_session_token = None6is_incremental_backup = 'N'

1def show_warn(warn_msg):2    """3    Display a warning message in a formatted HTML alert box.4
5    Parameters6    ----------7    warn_msg : str8        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
18def show_error(error_msg):19    """20    Display an error message in a formatted HTML alert box.21
22    Parameters23    ----------24    error_msg : str25        The error message to display.26    """27    display(HTML(f'''<div class="alert alert-block alert-danger">28    <b class="fa fa-solid fa-exclamation-triangle"></b>29    <div>30        <p><b>Error</b></p>31        <p>{error_msg}</p>32    </div>33</div>'''))34
35
36def show_success(success_msg):37    """38    Display a success message in a formatted HTML alert box.39
40    Parameters41    ----------42    success_msg : str43        The success message to display.44    """45    display(HTML(f'''<div class="alert alert-block alert-success">46    <b class="fa fa-solid fa-check-circle"></b>47    <div>48        <p><b>Success</b></p>49        <p>{success_msg}</p>50    </div>51</div>'''))

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

1def get_bkp_path(s3_path, db_name, do_increment):2    """3    Get the backup path based on the type of backup (incremental or initial).4
5    Parameters6    ----------7    s3_path : str8        The base S3 path for backups.9    db_name : str10        The name of the database.11    do_increment: str12        'Y' for incremental backup, else Full Backup13
14    Returns15    -------16    str17        The final backup path.18
19    """20    if do_increment == 'Y':21        logging.info('Is an incremental backup, will use exact path')22        return s3_path23    else:24        logging.info('Is an initial backup, will use time appended path')25        t = time.localtime(time.time())26        my_path = f'{s3_path}/{db_name}/{t.tm_year}-{t.tm_mon:02d}-{t.tm_mday:02d}/{t.tm_hour:02d}{t.tm_min:02d}{t.tm_sec:02d}/'27        logging.info(f'Backup Path : {my_path}')28        print(f'Backup Path : {my_path}')29        return my_path30
31
32
33def get_sql_statement(db_name_to_bkp, is_incremental_backup):34    """35    Get the SQL statement for backing up a database.36
37    Parameters38    ----------39    db_name_to_bkp : str40        The name of the database to backup.41    is_incremental_backup : str42        is incremental backup.43
44    Returns45    -------46    str47        The SQL statement for backup.48
49    """50    global aws_key_id, aws_secret_key, aws_region, s3_target_path, aws_session_token51    aws_key_id = (input('Enter AWS_API_KEY_ID:') if aws_key_id == None else aws_key_id)52    aws_secret_key = (getpass.getpass('Enter AWS_API_SECRET:') if aws_secret_key == None else aws_secret_key)53    aws_region = (input('Enter AWS_REGION:') if aws_region == None else aws_region)54    s3_target_path = (input('Enter AWS S3 Path:') if s3_target_path == None else s3_target_path)55    aws_session_token = (input('Enter AWS_SESSION_TOKEN:') if aws_session_token == None else aws_session_token)56
57    data = io.StringIO()58    data.write('BACKUP DATABASE ' + db_name_to_bkp + ' ')59    if is_incremental_backup == 'Y':60        data.write(' WITH DIFFERENTIAL ')61    else:62        data.write(' WITH INIT ')63    data.write(' TO S3 "' + get_bkp_path(s3_target_path, db_name_to_bkp, is_incremental_backup) + '" ')64    data.write(' CONFIG \'{"region":"' + aws_region + '"}\'')65    data.write(' CREDENTIALS \'{"aws_access_key_id":"' + aws_key_id66               + '","aws_secret_access_key":"' + aws_secret_key + '"' )67    if aws_session_token != '':68        data.write(', "aws_session_token":"' + aws_session_token + '" ')69
70    data.write('}\' ')71    logging.debug(f'statement: {data.getvalue()}')72    return data.getvalue()73
74
75def perform_backup(my_cursor, curr_db_name, do_incremental):76    """77    Perform a database backup.78
79    Parameters80    ----------81    my_cursor : cursor82        The database cursor.83    curr_db_name : str84        The name of the database to backup.85    do_incremental: str86        'Y' to perform incremental backup87
88    """89    logging.debug(f'backing up db {curr_db_name}')90    my_cursor.execute(get_sql_statement(curr_db_name, do_incremental))91    results = cursor.fetchall()92    if results is None:93        logging.error('Backup execution failed')94    else:95        logging.info("Backup completed")

1enable_debug_log(True)2print(connection_url)3try:4    if connection_url.endswith('/') or connection_url.endswith('information_schema'):5        #Hanlde case when database not selected or information_schema selected6        #Connect to information schema and backup all databases7        logging.debug('No database selected, will use information_schema and back up all databases')8        my_db_url = connection_url9        if connection_url.endswith('/'):10            my_db_url = my_db_url + 'information_schema'11            logging.debug(f'connection url updated {my_db_url}')12
13        conn = s2.connect(my_db_url, results_type='dict')14        with conn.cursor() as cursor:15            # Get a list of databases to backup16            cursor.execute(17                "SELECT schema_name FROM information_schema.schemata WHERE  schema_name NOT IN ( 'cluster', 'memsql', 'information_schema' );")18            for row in cursor.fetchall():19                logging.debug(f"processing db {row['schema_name']}")20                # Backup each database21                perform_backup(my_cursor=cursor, curr_db_name=row['schema_name'], do_incremental='N')22                logging.debug(f"processing db {row['schema_name']} complete")23
24    else:25            #Connect to selected database and take its backup26            database_to_bkp = connection_url.split('/')[-1]27
28            # Establish a connection to the database29            conn = s2.connect(results_type='dict')30            with conn.cursor() as cursor:31                perform_backup(my_cursor=cursor, curr_db_name=database_to_bkp, do_incremental=is_incremental_backup)32
33            show_success('Backup Process Completed')34
35except s2.exceptions.OperationalError as ope:36    # Handle specific operational errors37    if 'NoSuchBucket' in ope.errmsg:38        logging.error('Provided S3 Bucket does not exist. Please check.')39        show_error('Provided S3 Bucket does not exist. Please check.')40    elif 'Access denied' in ope.errmsg:41        logging.error('Failed to backup due to missing grants or firewall settings. Please check.')42        show_error('Failed to backup due to missing grants or firewall settings. Please check.')43    else:44        logging.error(f'Failed. Error message: {ope.errmsg}')45        show_error(f'Failed to backup. {ope.errmsg}')46except s2.Error as e:47    # Handle any other errors48    print(f'Encountered exception {e}')49    logging.exception(e)50    show_error(f'Failed to backup. {str(e)}')51
52print('\n\nScript execution completed')