feat(cli): support mv command

fix: parallel respects block_size
feat: lease command supports wait_sec

Author: william-silversmith

taskqueue/aws_queue_api.py

@@ -116,7 +116,10 @@ def cancel_lease(self, rhandle):
   def release_all(self):
     raise NotImplementedError()
 
-  def _request(self, num_tasks, visibility_timeout):
+  def lease(self, seconds, num_tasks=1, wait_sec=20):
+    if wait_sec is None:
+      wait_sec = 20
+
     resp = self.sqs.receive_message(
       QueueUrl=self.qurl,
       AttributeNames=[
@@ -126,8 +129,8 @@ def _request(self, num_tasks, visibility_timeout):
       MessageAttributeNames=[
         'All'
       ],
-      VisibilityTimeout=visibility_timeout,
-      WaitTimeSeconds=20,
+      VisibilityTimeout=seconds,
+      WaitTimeSeconds=wait_sec,
     )
 
     if 'Messages' not in resp:
@@ -140,11 +143,6 @@ def _request(self, num_tasks, visibility_timeout):
       tasks.append(task)
     return tasks
 
-  def lease(self, seconds, num_tasks=1):
-    if num_tasks > 1:
-      raise ValueError("This library (not boto/SQS) only supports fetching one task at a time. Requested: {}.".format(num_tasks))
-    return self._request(num_tasks, seconds)
-
   def delete(self, task):
     if type(task) == str:
       rhandle = task
@@ -175,13 +173,13 @@ def purge(self):
 
     while self.enqueued:
       # visibility_timeout must be > 0 for delete to work
-      tasks = self._request(num_tasks=10, visibility_timeout=10)
+      tasks = self.lease(num_tasks=10, seconds=10)
       for task in tasks:
         self.delete(task)
     return self
 
   def __iter__(self):
-    return iter(self._request(num_tasks=10, visibility_timeout=0))
+    return iter(self.lease(num_tasks=10, seconds=0))
 
 
 

taskqueue/file_queue_api.py

@@ -2,6 +2,7 @@
 import functools
 import itertools
 import json
+import math
 import operator
 import os.path
 import random
@@ -320,7 +321,10 @@ def _lease_filename(self, filename, seconds):
 
     return json.loads(read_file(new_filepath))
 
-  def lease(self, seconds, num_tasks):
+  def lease(self, seconds, num_tasks, wait_sec=None):
+    if wait_sec is None:
+      wait_sec = 0
+
     def fmt(direntry):
       filename = direntry.name
       timestamp, _ = filename.split('--')
@@ -354,7 +358,24 @@ def fmt(direntry):
       if lessee is not None:
         leases.append(lessee)
 
-    return leases
+    wait_leases = []
+    if wait_sec > 0 and len(leases) < num_tasks:
+      # Add a constant b/c this will cascade into shorter and 
+      # shorter checks as wait_sec shrinks and we don't
+      # want hundreds of workers to accidently synchronize
+      sleep_amt = random.random() * (wait_sec + 1)
+
+      # but we still want to guarantee that wait_sec is not
+      # exceeded.
+      sleep_amt = min(sleep_amt, wait_sec)
+      time.sleep(sleep_amt)
+      wait_leases = self.lease(
+        seconds, 
+        num_tasks - len(leases), 
+        wait_sec - sleep_amt
+      )
+
+    return leases + wait_leases
 
   @retry
   def delete(self, task):

taskqueue/taskqueue.py

@@ -183,7 +183,7 @@ def insert(
     )
 
     def insertfn(batch):
-        return self.api.insert(batch, delay_seconds) 
+      return self.api.insert(batch, delay_seconds) 
 
     cts = schedule_jobs(
       fns=( partial(insertfn, batch) for batch in sip(bodies, batch_size) ),
@@ -219,7 +219,7 @@ def cancel(self, task):
   def release_all(self):
     return self.api.release_all()
 
-  def lease(self, seconds=600, num_tasks=1):
+  def lease(self, seconds=600, num_tasks=1, wait_sec=None):
     """
     Acquires a lease on the topmost N unowned tasks in the specified queue.
     Required query parameters: leaseSecs, numTasks
@@ -229,7 +229,7 @@ def lease(self, seconds=600, num_tasks=1):
     if seconds < 0:
       raise ValueError("lease seconds must be >= 0. Got: " + str(seconds))
 
-    tasks = self.api.lease(seconds, num_tasks)
+    tasks = self.api.lease(seconds, num_tasks, wait_sec)
 
     if not len(tasks):
       raise QueueEmptyError()
@@ -562,7 +562,7 @@ def capturing_soloprocess_upload(*args, **kwargs):
   ct = 0
   with tqdm(desc="Upload", total=total) as pbar:
     with pathos.pools.ProcessPool(parallel) as pool:
-      for num_inserted in pool.imap(uploadfn, sip(tasks, 2000)):
+      for num_inserted in pool.imap(uploadfn, sip(tasks, block_size)):
         pbar.update(num_inserted)
         ct += num_inserted
 

taskqueue_cli/taskqueue_cli.py

@@ -2,7 +2,7 @@
 
 import click
 
-from taskqueue import TaskQueue, __version__
+from taskqueue import TaskQueue, __version__, QueueEmptyError
 from taskqueue.lib import toabs
 from taskqueue.paths import get_protocol
 
@@ -72,7 +72,8 @@ def cp(src, dest):
   process while a queue is being worked.
 
   Currently sqs queues are not copiable,
-  but you can copy an fq to sqs.
+  but you can copy an fq to sqs. The mv
+  command supports sqs queues.
   """
   src = normalize_path(src)
   dest = normalize_path(dest)
@@ -82,5 +83,36 @@ def cp(src, dest):
 
   tqd.insert(tqs)
 
+@main.command()
+@click.argument("src")
+@click.argument("dest")
+def mv(src, dest):
+  """
+  Moves the contents of a queue to another
+  service or location. Do not run this
+  process while a queue is being worked.
+
+  Moving an sqs queue to a file queue
+  may result in duplicated tasks.
+  """
+  src = normalize_path(src)
+  dest = normalize_path(dest)
+
+  tqd = TaskQueue(dest)
+  tqs = TaskQueue(src)
+
+  while True:
+    try:
+      tasks = tqs.lease(num_tasks=10, seconds=10)
+    except QueueEmptyError:
+      break
+
+    tqd.insert(tasks)
+    tqs.delete(tasks)
+
+
+
+
+
 
 

test/test_taskqueue.py

@@ -9,7 +9,11 @@
 import pytest
 
 import taskqueue
-from taskqueue import queueable, FunctionTask, RegisteredTask, TaskQueue, MockTask, PrintTask, LocalTaskQueue
+from taskqueue import (
+  queueable, FunctionTask, RegisteredTask, 
+  TaskQueue, MockTask, PrintTask, LocalTaskQueue,
+  QueueEmptyError
+)
 from taskqueue.paths import ExtractedPath, mkpath
 from taskqueue.queueables import totask
 from taskqueue.queueablefns import tofunc, UnregisteredFunctionError, func2task
@@ -154,6 +158,28 @@ def test_get(sqs, protocol):
     t = tq.lease()
     tq.delete(t)
 
+def test_lease(sqs):
+  path = getpath("sqs") 
+  tq = TaskQueue(path, n_threads=0)
+
+  n_inserts = 20
+  tq.purge()
+  tq.insert(( PrintTask(str(x)) for x in range(n_inserts) ))
+
+  tasks = tq.lease(num_tasks=10, wait_sec=0)
+  assert len(tasks) == 10
+  tq.delete(tasks)
+
+  tasks = tq.lease(num_tasks=10, wait_sec=0)
+  assert len(tasks) == 10
+  tq.delete(tasks)
+
+  try:
+    tasks = tq.lease(num_tasks=10, wait_sec=0)
+    assert False
+  except QueueEmptyError:
+    pass
+
 @pytest.mark.parametrize('protocol', PROTOCOL)
 def test_single_threaded_insertion(sqs, protocol):
   path = getpath(protocol)