Merge pull request #8 from andersinno/mpz-11-enqueue-objects-by-ids

deferred-enqueue-objects: Add support for object-ids parameter

Author: suutari-ai

src/passari_workflow/jobs/enqueue_objects.py

@@ -2,9 +2,9 @@
     enqueue_objects as do_enqueue_objects
 
 
-def enqueue_objects(object_count):
+def enqueue_objects(object_count=None, object_ids=None):
     """
     Enqueue objects using a background RQ job. This is used in when we don't
     want the user action to block.
     """
-    do_enqueue_objects(object_count)
+    do_enqueue_objects(object_count, object_ids)

src/passari_workflow/scripts/deferred_enqueue_objects.py

@@ -10,26 +10,51 @@
 from passari_workflow.jobs.enqueue_objects import enqueue_objects
 
 
-def deferred_enqueue_objects(object_count):
+def deferred_enqueue_objects(object_count=None, object_ids=None):
     """
     Enqueue given number of objects to the preservation workflow using a
     background RQ job
 
     :param int object_count: How many objects to enqueue at most
+    :param list object_ids: Specific object IDs to enqueue
     """
     queue = get_queue(QueueType.ENQUEUE_OBJECTS)
-    queue.enqueue(enqueue_objects, kwargs={"object_count": object_count})
 
-    print(f"{object_count} object(s) will be enqueued")
+    if object_ids:
+        queue.enqueue(enqueue_objects, kwargs={"object_ids": object_ids})
+        print(f"{len(object_ids)} object(s) with IDs {object_ids} will be enqueued")
+    else:
+        queue.enqueue(enqueue_objects, kwargs={"object_count": object_count})
+        print(f"{object_count} object(s) will be enqueued")
 
-    return object_count
+    return object_count or len(object_ids)
 
 
 @click.command()
 @click.option(
-    "--object-count", default=10, help="How many objects to enqueue")
-def cli(object_count):
-    deferred_enqueue_objects(object_count)
+    "--object-count", type=int, default=None,
+    help="How many objects to enqueue"
+)
+@click.option(
+    "--object-ids", type=str, default=None,
+    help="Comma-separated list of specific object IDs to enqueue, e.g., '1,2,3,4'"
+)
+def cli(object_count, object_ids):
+    if object_count and object_ids:
+        raise click.UsageError(
+            "You cannot use both --object-count and --object-ids at the same time."
+        )
+
+    if not object_count and not object_ids:
+        raise click.UsageError(
+            "You must provide either --object-count or --object-ids."
+        )
+
+    object_ids_list = object_ids.split(",") if object_ids else None
+    deferred_enqueue_objects(
+        object_count=object_count,
+        object_ids=object_ids_list
+    )
 
 
 if __name__ == "__name__":

src/passari_workflow/scripts/enqueue_objects.py

@@ -30,7 +30,7 @@ def enqueue_object(object_id):
     )
 
 
-def enqueue_objects(object_count, random=False, object_ids=None):
+def enqueue_objects(object_count, object_ids=None, random=False):
     """
     Enqueue given number of objects to the preservation workflow.
 

tests/scripts/test_deferred_enqueue_objects.py

@@ -37,3 +37,29 @@ def test_deferred_enqueue_objects(
     # 5 objects should now be enqueued
     queue = get_queue(QueueType.DOWNLOAD_OBJECT)
     assert len(queue.job_ids) == 5
+
+
+def test_deferred_enqueue_objects_by_ids(
+        redis, session, deferred_enqueue_objects, museum_object_factory):
+    """
+    Enqueue 3 objects by providing their ids
+    """
+    for i in range(0, 20):
+        museum_object_factory(
+            id=i, preserved=False,
+            metadata_hash="", attachment_metadata_hash=""
+        )
+
+    result = deferred_enqueue_objects(["--object-ids", "1,2,3"])
+    assert "3 object(s) with IDs ['1', '2', '3'] will be enqueued" in result.stdout
+
+    queue = get_queue(QueueType.ENQUEUE_OBJECTS)
+    # One job is enqueued
+    assert len(queue.job_ids) == 1
+
+    # Finish the job
+    SimpleWorker([queue], connection=queue.connection).work(burst=True)
+
+    # 3 objects should now be enqueued
+    queue = get_queue(QueueType.DOWNLOAD_OBJECT)
+    assert len(queue.job_ids) == 3