import time from boto.dynamodb.exceptions import DynamoDBResponseError from boto.dynamodb2.table import Table from boto.exception import JSONResponseError from garcon import task def _get_current_wcu_and_rcu(table): """ Get current values of WCU and RCU from DynamoDB table. We need time.sleep(10) as a first line, because invocation of the describe() method right after a plane request could provoke ThrottlingException. Args: table (boto.dynamodb2.Table object): The table object. Returns: tuple: current_wcu (int), current_rcu (int). """ time.sleep(10) # see the docstring above table_description = table.describe().get('Table') table_throughput = table_description.get( 'ProvisionedThroughput') current_rcu = table_throughput.get('ReadCapacityUnits') current_wcu = table_throughput.get('WriteCapacityUnits') time.sleep(10) return current_wcu, current_rcu def _is_throughput_changes_required( table, required_wcu=None, required_rcu=None): """ Check if the current values of WCU and RCU matches with required ones. Args: table (boto.dynamodb2.Table object): The table object. required_wcu (int, optional): Write capacity units. required_rcu (int, optional): Read capacity units. Returns: bool: True if we need to change throughput, False if not. """ current_wcu, current_rcu = _get_current_wcu_and_rcu(table) return (required_wcu and required_wcu != current_wcu) or ( required_rcu and required_rcu != current_rcu) def _sleep_if_table_is_busy(table, activity=None): """ Sleep until the table is ready. Args: table (boto.dynamodb2.Table object): The table object. activity (Garcon activity object, optional): The activity worker. Raises: DynamoDBResponseError: An exception with a custom message. """ time.sleep(10) statuses = ['UPDATING', 'CREATING', 'DELETING'] current_status = table.describe().get('Table').get( 'TableStatus') time.sleep(10) # 40 minutes, because changing throughput of big tables could take a lot retries_timeout = 60 * 40 while current_status in statuses: 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 )) if retries_timeout > 0: current_status = table.describe().get('Table').get( 'TableStatus') time.sleep(10) retries_timeout -= 10 else: # We should reach AWS if this exception appears constantly raise DynamoDBResponseError(400, 'Timeout', body={ '__type': 'com.amazon.coral.availability#ResourceInUseException', 'message': 'The Orchard сustom exception: update or delete or ' 'create table operation took more than 40 mins'}) def _update_table(table, activity=None, required_wcu=None, required_rcu=None): """ Update table with the new values. If only one of the values is provided, use the current one. Args: table (boto.dynamodb2.Table object): The table object. activity (ActivityWorker, optional): The SWF activity worker. required_wcu (int, optional): Write capacity units. required_rcu (int, optional): Read capacity units. """ current_wcu, current_rcu = _get_current_wcu_and_rcu( table) wcu = required_wcu or current_wcu rcu = required_rcu or current_rcu if activity: activity.logger.info( 'Updating {table} with WCU: {wcu} and RCU: {rcu}...'.format( table=table.table_name, wcu=wcu, rcu=rcu)) table.update(throughput=dict(write=wcu, read=rcu)) def update_table_params( table, activity=None, required_wcu=None, required_rcu=None): """ Update the throughput of a DynamoDB table. Args: table (boto.dynamodb2.Table object): The table object. activity (ActivityWorker, optional): The SWF activity worker. required_wcu (int, optional): Write capacity units. required_rcu (int, optional): Read capacity units. Returns (bool): True if throughput was updated, False if not. Raises: JSONResponseError. """ if not _is_throughput_changes_required(table, required_wcu, required_rcu): return False for retry_number in range(1, 6): try: _update_table(table, activity, required_wcu, required_rcu) except JSONResponseError as e: exception = e if e.error_code == 'ThrottlingException': if activity: activity.logger.info( 'An {exc} raises when performing throughput update ' 'during try {retry_number} of 5, trying to ' 'retry...'.format( exc=e.error_code, retry_number=retry_number)) time.sleep(10 * retry_number) _sleep_if_table_is_busy(table) if _is_throughput_changes_required( table, required_wcu, required_rcu): continue else: return True else: raise else: return True # If we are here, we have ThrottlingException even after 5 retries raise exception @task.decorate(timeout=3600) def set_table_throughput( activity, table_name, throughput_write=None, throughput_read=None): """Update the throughput of a DynamoDB table. 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_write (int, optional): The new throughput write value (WCU). throughput_read (int, optional): The new throughput read value (RCU). """ table = Table(table_name) result = update_table_params( table, activity, throughput_write, throughput_read) current_wcu, current_rcu = _get_current_wcu_and_rcu(table) if not throughput_write: throughput_write = current_wcu if not throughput_read: throughput_read = current_rcu if result: activity.logger.info( 'The throughtput of the DynamoDB table {table_name} was updated ' '(RCU: {rcu}, WCU: {wcu})'.format( table_name=table_name, rcu=throughput_read, wcu=throughput_write)) else: activity.logger.info( 'The throughtput of the DynamoDB table {table_name} already ' 'matches the required values (RCU: {rcu}, WCU: {wcu}' ')'.format( table_name=table_name, rcu=throughput_read, wcu=throughput_write)) @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. """ table = Table(table_name) _sleep_if_table_is_busy(table, activity=activity)