in source/client/python/api-v0.1/api/state_table_dynamodb.py [0:0]
def claim_task_for_agent(self, task_id, queue_handle_id, agent_id, expiration_timestamp):
"""
Once task has been consumed from the task queue, we need to update the owner and the state of the task in the state table.
Approximately equivalent to SQL:
Alter table state_table where TaskId == task_id
set Ownder = SelfWorkerID and task_status = Running and
condition to status == Pending and OwnerID == None
Args:
task_id: picked task
queue_handle_id: handle to manipulate visibility timeout at a later stage
agent_id: i.e., the worker ID that is working on this task
expiration_timestamp: heartbeat for TTL
Returns:
True if claim succesfull
Throws:
StateTableException on throttling
StateTableException on condition
Exception for all other errors
"""
session_id = self.__get_session_id_from_task_id(task_id)
claim_is_successful = True
try:
self.state_table.update_item(
Key={
'task_id': task_id
},
UpdateExpression="SET #var_task_owner = :val1, #var_task_status = :val2, #var_heartbeat_expiration_timestamp = :val3, #var_sqs_handler_id = :val4",
ExpressionAttributeValues={
':val1': agent_id,
':val2': self.__make_task_state_from_session_id(TASK_STATE_PROCESSING, session_id),
':val3': expiration_timestamp,
':val4': queue_handle_id
},
ExpressionAttributeNames={
"#var_task_owner": "task_owner",
"#var_task_status": "task_status",
"#var_heartbeat_expiration_timestamp": "heartbeat_expiration_timestamp",
"#var_sqs_handler_id": "sqs_handler_id"
},
ConditionExpression=Key('task_status').eq(
self.__make_task_state_from_session_id(TASK_STATE_PENDING, session_id)
) & Key('task_owner').eq('None')
)
except ClientError as e:
if e.response['Error']['Code'] == 'ConditionalCheckFailedException':
msg = f"Could not acquire task [{task_id}] for status [{self.__make_task_state_from_session_id(TASK_STATE_PENDING, session_id)}] from DynamoDB, someone else already locked it? [{e}]"
logging.warning(msg)
raise StateTableException(e, msg, caused_by_condition=True)
elif e.response['Error']['Code'] in ["ThrottlingException", "ProvisionedThroughputExceededException"]:
msg = f"Could not acquire task [{task_id}] from DynamoDB, Throttling Exception {e}"
logging.warning(msg)
raise StateTableException(e, msg, caused_by_throttling=True)
else:
msg = f"ClientError while acquire task [{task_id}] from DynamoDB: {e}"
logging.error(msg)
raise Exception(e)
except Exception as e:
msg = f"Failed to acquire task [{task_id}] for agent [{agent_id}]: from DynamoDB: {e}"
logging.error(msg)
raise e
return claim_is_successful