import time import boto3 from botocore.exceptions import ClientError from garcon import task class DynamoDBResponseError(ClientError): pass class DynamoUpdateError(Exception): pass def _sleep_if_table_is_busy(table_name, activity=None): """ Sleep until the table is ready. Args: table_name (str): table name. activity (Garcon activity object, optional): The activity worker. Raises: DynamoDBResponseError: An exception with a custom message. """ dynamodb = boto3.resource('dynamodb') statuses = ['UPDATING', 'CREATING', 'DELETING'] # 40 minutes, because changing throughput of big tables could take a lot retries_timeout = 40 * 60 attempt_sleep = 5 number_of_attempts = retries_timeout // attempt_sleep for attempt in range(number_of_attempts): table = dynamodb.Table(table_name) current_status = table.table_status if current_status not in statuses: return table if activity: activity.logger.info( 'Waiting for the DynamoDB table {table_name} to be ready ' '(current status: {status})'.format( table_name=table.table_name, status=current_status)) activity.heartbeat( details='Table %(table_name)s is not ready yet' % dict( table_name=table.table_name )) time.sleep(attempt_sleep) continue raise DynamoDBResponseError( error_response={ 'Error': { 'Code': 'com.amazon.coral.availability' '#ResourceInUseException', 'Message': '400 (Timeout): The Orchard ' 'сustom exception: update or delete or ' 'create table operation took more than 40 minutes'} }, operation_name='UpdateTable') def _update_table(table, throughput_settings, activity=None): """ Update table with the new values. Args: table (dynamodb.Table object): The table object. throughput_settings (dict): Throughput values for the table and global secondary indexes. activity (ActivityWorker, optional): The SWF activity worker. """ if activity: activity.logger.info( 'Updating {table} with following settings {settings}'.format( table=table.table_name, settings=throughput_settings)) table.update(**throughput_settings) def _update_table_params(table, throughput_settings, activity=None): """ Update the throughput of a DynamoDB table. Args: table (dynamodb.Table object): The table object. throughput_settings (dict): Throughput values for the table and global secondary indexes. activity (ActivityWorker, optional): The SWF activity worker. Returns (bool): True if throughput was updated, False if not. Raises: ClientError """ if not throughput_settings: return False for retry_number in range(1, 6): try: _update_table(table, throughput_settings, activity) except ClientError as e: if not e.response['Error']['Code'] == 'ThrottlingException': raise else: if activity: activity.logger.info( 'An {exc} raises when performing throughput update ' 'during try {retry_number} of 5, trying to ' 'retry...'.format( exc=e.response['Error']['Code'], retry_number=retry_number)) time.sleep(10 * retry_number) _sleep_if_table_is_busy(table) # If the table still hasn't been updated if _settings_for_update(table, throughput_settings): continue else: return True else: return True raise DynamoUpdateError('Still failed after all retries') def _table_current_settings(table): """Get current throughput values for the table. Args: table (dynamodb.Table object): The table object. Returns: dict: Throughput values of the table and all global secondary indexes. """ current_throughput = { 'throughput': { 'read': table.provisioned_throughput[ 'ReadCapacityUnits'], 'write': table.provisioned_throughput[ 'WriteCapacityUnits']} } global_secondary_indexes = table.global_secondary_indexes if global_secondary_indexes: current_throughput['global_indexes'] = {} for index in global_secondary_indexes: current_throughput['global_indexes'][index['IndexName']] = { 'read': index['ProvisionedThroughput']['ReadCapacityUnits'], 'write': index['ProvisionedThroughput']['WriteCapacityUnits']} return current_throughput def _settings_for_update(table, new_settings): """Compare new settings with current values and return only necessary. Args: table (dynamodb.Table object): The table object. new_settings (dict): Throughput values for the table and global secondary indexes. Returns: dict: Nesessary for update throughput settings. """ current_settings = _table_current_settings(table) necessary_settings = {} if 'throughput' in new_settings: n_read = new_settings['throughput'].get('read') c_read = current_settings['throughput'].get('read') n_write = new_settings['throughput'].get('write') c_write = current_settings['throughput'].get('write') if (n_read and n_read != c_read or n_write and n_write != c_write): necessary_settings['ProvisionedThroughput'] = { 'ReadCapacityUnits': n_read or c_read, 'WriteCapacityUnits': n_write or c_write } global_secondary_indexes = new_settings.get('global_indexes') if global_secondary_indexes: index_updates = [] for index_name in global_secondary_indexes: new_global_indexes = new_settings['global_indexes'][index_name] current_global_indexes = ( current_settings['global_indexes'][index_name]) n_read = new_global_indexes.get('read') c_read = current_global_indexes.get('read') n_write = new_global_indexes.get('write') c_write = current_global_indexes.get('write') if (n_read and n_read != c_read or n_write and n_write != c_write): index_updates.append({ 'Update': { 'IndexName': index_name, 'ProvisionedThroughput': { 'ReadCapacityUnits': n_read or c_read, 'WriteCapacityUnits': n_write or c_write}}}) if index_updates: necessary_settings['GlobalSecondaryIndexUpdates'] = index_updates return necessary_settings @task.decorate(timeout=3600) def set_table_throughput(activity, table_name, throughput_settings): """Update the throughput of a DynamoDB table and GSIs. Throughput updates are limited, so the task checks first the current table throughput before requesting a modification. If the throughput is the one desired, no request is sent. Args: activity (ActivityWorker): The SWF activity worker. table_name (str): The name of the table. throughput_settings (dict): Throughput values for the table and global secondary indexes in the following format: todo: shall we update the throughput_settings parameter format todo: accordingly to a new boto3 request/response JSON-type? { 'throughput': { 'read': 2, 'write': 3 }, 'global_indexes': { 'my-gsi-index-name': { 'read': 3, 'write': 1 } } """ dynamodb = boto3.resource('dynamodb') table = dynamodb.Table(table_name) settings_for_update = _settings_for_update(table, throughput_settings) result = _update_table_params(table, settings_for_update, activity) if result: activity.logger.info( 'The throughput of the DynamoDB table {table_name} was updated. ' 'New values: {settings}'.format( table_name=table_name, settings=settings_for_update)) else: activity.logger.info( 'The throughput of the DynamoDB table {table_name} already ' 'matches the required values: {settings}' ')'.format( table_name=table_name, settings=settings_for_update)) @task.decorate(timeout=3600) def wait_for_table_task_completion(activity, table_name): """Wait for the table task to complete. Actions such as updating, creating or deleting a table have delays (it can take a few seconds to several minutes.) This task wait until the table is ready before releasing the task. Args: activity (ActivityWorker): The SWF activity worker. table_name (string): The name of the table. """ _sleep_if_table_is_busy(table_name, activity=activity)