Locking; implement lock and unlock methods

s3pypi/locking.py

@@ -1,29 +1,36 @@
 import abc
+import datetime as dt
+import time
 from contextlib import contextmanager
 
 import boto3
 
 
+class LockTimeoutError(Exception):
+    def __init__(self, lock_id: str):
+        super().__init__(f"Timed out trying to acquire lock '{lock_id}'")
+
+
 class Locker(abc.ABC):
     @contextmanager
-    def __call__(self, key: str):
-        self._lock(key)
+    def __call__(self, lock_id: str):
+        self._lock(lock_id)
         try:
             yield
         finally:
-            self._unlock(key)
+            self._unlock(lock_id)
 
     @abc.abstractmethod
-    def _lock(self, key: str):
+    def _lock(self, lock_id: str):
         ...
 
     @abc.abstractmethod
-    def _unlock(self, key: str):
+    def _unlock(self, lock_id: str):
         ...
 
 
 class DummyLocker(Locker):
-    def _lock(self, key: str):
+    def _lock(self, lock_id: str):
         pass
 
     _unlock = _lock
@@ -34,14 +41,29 @@ def __init__(
         self,
         session: boto3.session.Session,
         table: str,
-        timeout: int = 10,
+        poll_interval: int = 1,
+        max_attempts: int = 30,
     ):
         db = session.resource("dynamodb")
         self.table = db.Table(table)
-        self.timeout = timeout
+        self.exc = self.table.meta.client.exceptions
+        self.poll_interval = poll_interval
+        self.max_attempts = max_attempts
+
+    def _lock(self, lock_id: str):
+        for attempt in range(1, self.max_attempts + 1):
+            now = dt.datetime.now(dt.timezone.utc)
+            try:
+                self.table.put_item(
+                    Item={"LockID": lock_id, "AcquiredAt": now.isoformat()},
+                    ConditionExpression="attribute_not_exists(LockID)",
+                )
+                return
+            except self.exc.ConditionalCheckFailedException:
+                if attempt < self.max_attempts:
+                    time.sleep(self.poll_interval)
 
-    def _lock(self, key: str):
-        raise NotImplementedError
+        raise LockTimeoutError(lock_id)
 
-    def _unlock(self, key: str):
-        raise NotImplementedError
+    def _unlock(self, lock_id: str):
+        self.table.delete_item(Item={"LockID": lock_id})