#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights served.
SPDX-License-Identifier: MIT-0
Permission is hereby granted, free of charge, to any person taining a copy of this
software and associated documentation files (the oftware"), to deal in the Software
without restriction, including without limitation the rights use, copy, modify,
merge, publish, distribute, sublicense, and/or sell copies the Software, and to
permit persons to whom the Software is furnished to do so.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY ND, EXPRESS OR IMPLIED,
INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF RCHANTABILITY, FITNESS FOR A
PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL E AUTHORS OR COPYRIGHT
HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, ETHER IN AN ACTION
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN NNECTION WITH THE
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""
from botocore.exceptions import ClientError
import logging
from commons import current_milli_time
from config_provider import config_provider
# Logger
logger = logging.getLogger(__name__)
logger.setLevel(config_provider.logging_level)
class DDBLeaseManager(object):
"""
Manages lease needed to start Poller Worker for Stream.
A Lease stores information about worker currently reading data from Stream and
also keep track of the records that have already been processed from Stream using Checkpoint.
A worker needs to take lease at beginning of the process and needs to evict lease at the
end of process or system failure. Two simultaneous workers can not take lease for Same Application
which is guaranteed by Conditional Check on Lease Owner.
Currently, Lease Information is stored in Dynamo DB Table which makes it simpler to manage taking lease, evicting
lease and updating lease process through Dynamo DB Conditional Check.
Note: Neptune stream is single sharded and will be consumed by single worker.
"""
def __init__(self, lease_table):
self.table = lease_table
def create_lease_if_not_exists(self, item_dict=None):
"""
Creates a new lease. Conditional on a lease not already existing.
"""
try:
self.table.put_item(Item=item_dict,
ConditionExpression='attribute_not_exists(leaseKey)')
logger.debug("Successfully Created Lease Entry - {}".format(str(item_dict)))
except ClientError as e:
if e.response.get('Error', {}).get('Code') != 'ConditionalCheckFailedException':
logger.error("Failed to put item in to {0} : error{1}".format(self.table, e))
raise
def get_lease(self, lease_key):
"""
Returns lease for a given key.
"""
try:
response = self.table.get_item(Key={'leaseKey': lease_key}, ConsistentRead=True)
return response['Item']
except Exception as e:
logger.error('{0} failed query items , error{1}'.format(self.table, e))
return None
def take_lease(self, item_dict=None):
"""
Take a lease for the given owner update its owner field. Conditional on
the leaseOwner is nobody in DynamoDB. Mutates the owner of the
passed-in lease object after updating DynamoDB.
"""
try:
response = self.table.update_item(
Key={'leaseKey': item_dict['leaseKey']},
UpdateExpression="SET leaseOwner = :leaseOwnerVal," +
" lastUpdateTime= :lastUpdateTimeVal",
ConditionExpression="leaseOwner = :NOBODY",
ExpressionAttributeValues={':leaseOwnerVal': item_dict['leaseOwner'],
':lastUpdateTimeVal': current_milli_time(),
':NOBODY': 'nobody'
},
ReturnValues="ALL_NEW")
logger.debug("Successfully Taken Lease for leaseKey - {}, leaseOwner - {}. Lease - {}"
.format(str(item_dict['leaseKey']), str(item_dict['leaseOwner']), str(response["Attributes"])))
return response["Attributes"]
except ClientError as e:
if e.response['Error']['Code'] == "ConditionalCheckFailedException":
logger.error(e.response['Error']['Message'])
else:
raise
def update_lease(self, item_dict=None):
"""
Update a lease by updating Checkpoint for Lease. Conditional on
the leaseOwner same as Lease owner at the time of taking lease in DynamoDB .
Checkpoint is the last processed event (Neptune Stream Checkpoint is consist of
checkpoint and checkpointSubSequenceNumber)
"""
try:
response = self.table.update_item(
Key={'leaseKey': item_dict['leaseKey']},
UpdateExpression="SET checkpoint = :checkpointVal," +
" checkpointSubSequenceNumber = :checkpointSubSequenceNumberVal," +
" lastUpdateTime= :lastUpdateTimeVal",
ConditionExpression="leaseOwner = :leaseOwnerVal ",
ExpressionAttributeValues={':leaseOwnerVal': item_dict['leaseOwner'],
':checkpointVal': item_dict['checkpoint'],
':checkpointSubSequenceNumberVal': item_dict['checkpointSubSequenceNumber'],
':lastUpdateTimeVal': current_milli_time()
},
ReturnValues="ALL_NEW")
logger.debug("Successfully Updated Lease - {}".format(str(response["Attributes"])))
return response["Attributes"]
except ClientError as e:
if e.response['Error']['Code'] == "ConditionalCheckFailedException":
logger.error(e.response['Error']['Message'])
else:
raise
def evict_lease(self, item_dict=None):
"""
Evict the current owner of lease by setting owner to nobody. Conditional on the owner in DynamoDB matching
the owner of the input. Mutates the owner of the passed-in lease object after updating
the record in DynamoDB.
"""
try:
response = self.table.update_item(
Key={'leaseKey': item_dict['leaseKey']},
UpdateExpression="SET leaseOwner = :NOBODY," +
" lastUpdateTime= :lastUpdateTimeVal",
ConditionExpression="leaseOwner = :leaseOwnerVal",
ExpressionAttributeValues={':lastUpdateTimeVal': current_milli_time(),
':NOBODY': 'nobody',
':leaseOwnerVal': item_dict['leaseOwner']
},
ReturnValues="ALL_NEW")
logger.debug("Successfully Evicted Lease - {}".format(str(response["Attributes"])))
return response["Attributes"]
except ClientError as e:
if e.response['Error']['Code'] == "ConditionalCheckFailedException":
logger.error(e.response['Error']['Message'])
else:
raise
def delete_all_items_in_lease_table(self):
"""
Delete all lease items in dynamo db table
"""
response = self.table.scan()
items = response['Items']
number_of_items = len(items)
if number_of_items == 0: # no items to delete
return
for item in items:
self.table.delete_item(Key={'leaseKey': item['leaseKey']})