feat: implement renew_lease for sqs

seung-lab · Sep 1, 2023 · f164cd0 · f164cd0
1 parent 1d13d60
commit f164cd0
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/taskqueue/aws_queue_api.py b/taskqueue/aws_queue_api.py
@@ -143,11 +143,16 @@ def add_insert_count(self, ct):
   def rezero(self):
     pass
 
-  def renew_lease(self, seconds):
-    raise NotImplementedError() 
+  @retry
+  def renew_lease(self, task, seconds):
+    self.sqs.change_message_visibility(
+      QueueUrl=self.qurl,
+      ReceiptHandle=task.id,
+      VisibilityTimeout=seconds,
+    )
 
-  def cancel_lease(self, rhandle):
-    raise NotImplementedError()
+  def cancel_lease(self, task):
+    self.renew_lease(task, 0)
 
   def release_all(self):
     raise NotImplementedError()