-
-
Save gregeinfrank/7fe8d0e4aa7459d635de6f5de30a5721 to your computer and use it in GitHub Desktop.
A lock implemented in AWS DynamoDB
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import boto3 | |
import datetime | |
import logging | |
import sys | |
import traceback | |
from contextlib import contextmanager | |
logger = logging.getLogger() | |
class DynamoDBLock(object): | |
def __init__(self, lock_name, dynamodb_table_name, **session_kwargs): | |
self.lock_name = lock_name | |
self.should_release_lock_when_done = True | |
self.dynamodb_table_name = dynamodb_table_name | |
self.dynamodb = boto3.session.Session(**session_kwargs).client('dynamodb') | |
@contextmanager | |
def lock(self): | |
self._obtain_lock() | |
try: | |
yield | |
self.release_lock() | |
except Exception as e: | |
# Explicitly *not* releasing the lock when there is a failure, until an | |
# engineer can fix the root issue and manually release the lock. | |
# To have the lock released on erors, add a `finally` block and call | |
# `release_lock` from there. | |
logger.error("Exception {} {} raised, not releasing lock {}".format( | |
e, e.message, self.lock_name)) | |
traceback.print_exc() | |
def _obtain_lock(self): | |
try: | |
self.dynamodb.put_item( | |
TableName=self.dynamodb_table_name, | |
Item={ | |
'LockType': {'S': self.lock_name}, | |
'LockedAt': {'S': datetime.datetime.utcnow().isoformat()}, | |
}, | |
# This will fail if the lock cannot be obtained | |
ConditionExpression='attribute_not_exists(LockType)' | |
) | |
except boto3.exceptions.botocore.exceptions.ClientError as e: | |
err_msg = e.response['Error']['Message'] | |
e.response['Error']['Message'] = "{} - unable to obtain {} lock".format( | |
err_msg, self.lock_name) | |
raise type(e), type(e)(e.response, e.operation_name), sys.exc_info()[2] | |
def release_lock(self): | |
if self.should_release_lock_when_done: | |
try: | |
self.dynamodb.delete_item( | |
TableName=self.dynamodb_table_name, | |
Key={'LockType': {'S': self.lock_name}} | |
) | |
except boto3.exceptions.botocore.exceptions.ClientError as e: | |
err_msg = e.response['Error']['Message'] | |
e.response['Error']['Message'] = "{} - unable to release {} lock".format( | |
err_msg, self.lock_name) | |
raise type(e), type(e)(e.response, e.operation_name), sys.exc_info()[2] |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment