feat: add purge command to cli, add status support to SQS

Author: william-silversmith

taskqueue/aws_queue_api.py

@@ -57,15 +57,16 @@ def enqueued(self):
 
   @property
   def inserted(self):
-    raise NotImplementedError()
+    return float('NaN')
 
   @property
   def completed(self):
-    raise NotImplementedError()
+    return float('NaN')
 
   @property
   def leased(self):
-    raise NotImplementedError()
+    status = self.status()
+    return int(status['ApproximateNumberOfMessagesNotVisible'])
 
   def is_empty():
     return self.enqueued == 0
@@ -165,18 +166,18 @@ def delete(self, task):
   def tally(self):
     pass
 
-  def purge(self):
-    # This is more efficient, but it kept freezing
-    # try:
-    #     self.sqs.purge_queue(QueueUrl=self.qurl)
+  def purge(self, native=False):
+    # can throw:
     # except botocore.errorfactory.PurgeQueueInProgress:
+    if native:
+      self.sqs.purge_queue(QueueUrl=self.qurl)
+      return
 
     while self.enqueued:
       # visibility_timeout must be > 0 for delete to work
       tasks = self.lease(num_tasks=10, seconds=10)
       for task in tasks:
         self.delete(task)
-    return self
 
   def __iter__(self):
     return iter(self.lease(num_tasks=10, seconds=0))

taskqueue/file_queue_api.py

@@ -416,7 +416,9 @@ def tally(self):
     with open(self.completions_path, 'ba') as f:
       f.write(b'\0')
 
-  def purge(self):
+  def purge(self, native=False):
+    # native only has meaning for SQS for now
+    # but we have to accept it as a parameter.
     all_files = itertools.chain(
       os.scandir(self.queue_path),
       os.scandir(self.movement_path)
@@ -427,6 +429,8 @@ def purge(self):
       except FileNotFoundError:
         pass
 
+    self.rezero()
+
   def is_empty(self):
     try:
       first(os.scandir(self.queue_path))

taskqueue/taskqueue.py

@@ -258,10 +258,10 @@ def deltask(tid):
       green=self.green,
     )
 
-  def purge(self):
+  def purge(self, native=False):
     """Deletes all tasks in the queue."""
     try:
-      return self.api.purge()
+      return self.api.purge(native)
     except AttributeError:
       while True:
         lst = self.list()

taskqueue_cli/taskqueue_cli.py

@@ -1,4 +1,5 @@
 import os
+import math
 
 import click
 from tqdm import tqdm
@@ -44,13 +45,18 @@ def status(queuepath):
   enq = tq.enqueued
   comp = tq.completed
   leased = tq.leased
-  print(f"Inserted: {ins}")
+
+  if not math.isnan(ins):
+    print(f"Inserted: {ins}")
+
   if ins > 0:
     print(f"Enqueued: {enq} ({enq / ins * 100:.1f}% left)")
-    print(f"Completed: {comp} ({comp / ins * 100:.1f}%)")
+    if not math.isnan(comp):
+      print(f"Completed: {comp} ({comp / ins * 100:.1f}%)")
   else:
     print(f"Enqueued: {enq} (--% left)")
-    print(f"Completed: {comp} (--%)")
+    if not math.isnan(comp):
+      print(f"Completed: {comp} (--%)")
 
   if enq > 0:
     print(f"Leased: {leased} ({leased / enq * 100:.1f}% of queue)")
@@ -118,8 +124,13 @@ def mv(src, dest):
       tqs.delete(tasks)
       pbar.update(len(tasks))
 
-
-
+@main.command()
+@click.argument("queuepath")
+def purge(queuepath):
+  """Delete all queued messages and zero out queue statistics."""
+  queuepath = normalize_path(queuepath)
+  tq = TaskQueue(queuepath)
+  tq.purge()