backpopulate files by resource id

learning_resources/management/commands/backpopulate_mit_edx_files.py

@@ -28,11 +28,25 @@ def add_arguments(self, parser):
             help="Overwrite any existing records",
         )
 
+        parser.add_argument(
+            "--resource-ids",
+            dest="learning_resource_ids",
+            required=False,
+            help="If set, backpopulate only the learning resources with these ids",
+        )
+
     def handle(self, *args, **options):  # noqa: ARG002
         """Run Populate MIT edX course run files"""
         chunk_size = options["chunk_size"]
+        resource_ids = (
+            options["learning_resource_ids"].split(",")
+            if options["learning_resource_ids"]
+            else None
+        )
         task = import_all_mit_edx_files.delay(
-            chunk_size=chunk_size, overwrite=options["force_overwrite"]
+            chunk_size=chunk_size,
+            overwrite=options["force_overwrite"],
+            learning_resource_ids=resource_ids,
         )
         self.stdout.write(f"Started task {task} to get MIT edX course run file data")
         self.stdout.write("Waiting on task...")

learning_resources/management/commands/backpopulate_mitxonline_files.py

@@ -28,11 +28,25 @@ def add_arguments(self, parser):
             help="Overwrite any existing records",
         )
 
+        parser.add_argument(
+            "--resource-ids",
+            dest="learning_resource_ids",
+            required=False,
+            help="If set, backpopulate only the learning resources with these ids",
+        )
+
     def handle(self, *args, **options):  # noqa: ARG002
         """Run Populate MITX Online course run files"""
         chunk_size = options["chunk_size"]
+        resource_ids = (
+            options["learning_resource_ids"].split(",")
+            if options["learning_resource_ids"]
+            else None
+        )
         task = import_all_mitxonline_files.delay(
-            chunk_size=chunk_size, overwrite=options["force_overwrite"]
+            chunk_size=chunk_size,
+            overwrite=options["force_overwrite"],
+            learning_resource_ids=resource_ids,
         )
         self.stdout.write(
             f"Started task {task} to get MITX Online course run file data"

learning_resources/management/commands/backpopulate_oll_files.py

@@ -28,14 +28,28 @@ def add_arguments(self, parser):
             help="Overwrite any existing records",
         )
 
+        parser.add_argument(
+            "--resource-ids",
+            dest="learning_resource_ids",
+            required=False,
+            help="If set, backpopulate only the learning resources with these ids",
+        )
+
     def handle(self, *args, **options):  # noqa: ARG002
         """Run Populate OLL course run files"""
         if not settings.OLL_LEARNING_COURSE_BUCKET_NAME:
             self.stderr.write("OLL contentfile settings not configured, skipping")
             return
         chunk_size = options["chunk_size"]
+        resource_ids = (
+            options["learning_resource_ids"].split(",")
+            if options["learning_resource_ids"]
+            else None
+        )
         task = import_all_oll_files.delay(
-            chunk_size=chunk_size, overwrite=options["force_overwrite"]
+            chunk_size=chunk_size,
+            overwrite=options["force_overwrite"],
+            learning_resource_ids=resource_ids,
         )
         self.stdout.write(f"Started task {task} to get OLL course run file data")
         self.stdout.write("Waiting on task...")

learning_resources/management/commands/backpopulate_xpro_files.py

@@ -27,12 +27,25 @@ def add_arguments(self, parser):
             action="store_true",
             help="Overwrite any existing records",
         )
+        parser.add_argument(
+            "--resource-ids",
+            dest="learning_resource_ids",
+            required=False,
+            help="If set, backpopulate only the learning resources with these ids",
+        )
 
     def handle(self, *args, **options):  # noqa: ARG002
         """Run Populate xpro course run files"""
         chunk_size = options["chunk_size"]
+        resource_ids = (
+            options["learning_resource_ids"].split(",")
+            if options["learning_resource_ids"]
+            else None
+        )
         task = import_all_xpro_files.delay(
-            chunk_size=chunk_size, overwrite=options["force_overwrite"]
+            chunk_size=chunk_size,
+            overwrite=options["force_overwrite"],
+            learning_resource_ids=resource_ids,
         )
         self.stdout.write(f"Started task {task} to get xpro course run file data")
         self.stdout.write("Waiting on task...")

learning_resources/tasks.py

@@ -141,13 +141,15 @@ def get_content_files(
     clear_search_cache()
 
 
-def get_content_tasks(
+def get_content_tasks(  # noqa: PLR0913
     etl_source: str,
     *,
     chunk_size: int | None = None,
     s3_prefix: str | None = None,
     override_base_prefix: bool = False,
     overwrite: bool = False,
+    learning_resource_ids: list[int] | None = None,
+    # Updated parameter
 ) -> celery.group:
     """
     Return a list of grouped celery tasks for indexing edx content
@@ -159,39 +161,54 @@ def get_content_tasks(
     archive_keys = get_most_recent_course_archives(
         etl_source, s3_prefix=s3_prefix, override_base_prefix=override_base_prefix
     )
+
+    if learning_resource_ids:
+        learning_resources = LearningResource.objects.filter(
+            id__in=learning_resource_ids, etl_source=etl_source
+        ).values_list("id", flat=True)
+    else:
+        learning_resources = (
+            LearningResource.objects.filter(
+                published=True, course__isnull=False, etl_source=etl_source
+            )
+            .exclude(readable_id__in=blocklisted_ids)
+            .order_by("-id")
+            .values_list("id", flat=True)
+        )
+
     return celery.group(
         [
             get_content_files.si(
                 ids, etl_source, archive_keys, s3_prefix=s3_prefix, overwrite=overwrite
             )
             for ids in chunks(
-                LearningResource.objects.filter(
-                    published=True, course__isnull=False, etl_source=etl_source
-                )
-                .exclude(readable_id__in=blocklisted_ids)
-                .order_by("-id")
-                .values_list("id", flat=True),
+                learning_resources,
                 chunk_size=chunk_size,
             )
         ]
     )
 
 
 @app.task(bind=True)
-def import_all_mit_edx_files(self, *, chunk_size=None, overwrite=False):
+def import_all_mit_edx_files(
+    self, *, chunk_size=None, overwrite=False, learning_resource_ids=None
+):
     """Ingest MIT edX files from an S3 bucket"""
     return self.replace(
         get_content_tasks(
             ETLSource.mit_edx.name,
             chunk_size=chunk_size,
             s3_prefix=settings.EDX_LEARNING_COURSE_BUCKET_PREFIX,
             overwrite=overwrite,
+            learning_resource_ids=learning_resource_ids,
         )
     )
 
 
 @app.task(bind=True)
-def import_all_oll_files(self, *, chunk_size=None, overwrite=False):
+def import_all_oll_files(
+    self, *, chunk_size=None, overwrite=False, learning_resource_ids=None
+):
     """Ingest MIT edX files from an S3 bucket"""
     return self.replace(
         get_content_tasks(
@@ -200,27 +217,39 @@ def import_all_oll_files(self, *, chunk_size=None, overwrite=False):
             s3_prefix=settings.OLL_LEARNING_COURSE_BUCKET_PREFIX,
             override_base_prefix=True,
             overwrite=overwrite,
+            learning_resource_ids=learning_resource_ids,
         )
     )
 
 
 @app.task(bind=True)
-def import_all_mitxonline_files(self, *, chunk_size=None, overwrite=False):
+def import_all_mitxonline_files(
+    self, *, chunk_size=None, overwrite=False, learning_resource_ids=None
+):
     """Ingest MITx Online files from an S3 bucket"""
+
     return self.replace(
         get_content_tasks(
-            ETLSource.mitxonline.name, chunk_size=chunk_size, overwrite=overwrite
+            ETLSource.mitxonline.name,
+            chunk_size=chunk_size,
+            overwrite=overwrite,
+            learning_resource_ids=learning_resource_ids,
         )
     )
 
 
 @app.task(bind=True)
-def import_all_xpro_files(self, *, chunk_size=None, overwrite=False):
+def import_all_xpro_files(
+    self, *, chunk_size=None, overwrite=False, learning_resource_ids=None
+):
     """Ingest xPRO OLX files from an S3 bucket"""
 
     return self.replace(
         get_content_tasks(
-            ETLSource.xpro.name, chunk_size=chunk_size, overwrite=overwrite
+            ETLSource.xpro.name,
+            chunk_size=chunk_size,
+            overwrite=overwrite,
+            learning_resource_ids=learning_resource_ids,
         )
     )
 

learning_resources/tasks_test.py

@@ -145,12 +145,15 @@ def test_import_all_mit_edx_files(settings, mocker, mocked_celery, mock_blocklis
         "learning_resources.tasks.get_content_tasks", autospec=True
     )
     with pytest.raises(mocked_celery.replace_exception_class):
-        tasks.import_all_mit_edx_files.delay(chunk_size=4, overwrite=False)
+        tasks.import_all_mit_edx_files.delay(
+            chunk_size=4, overwrite=False, learning_resource_ids=[1]
+        )
     get_content_tasks_mock.assert_called_once_with(
         ETLSource.mit_edx.name,
         chunk_size=4,
         s3_prefix="simeon-mitx-course-tarballs",
         overwrite=False,
+        learning_resource_ids=[1],
     )
 
 
@@ -163,9 +166,14 @@ def test_import_all_mitxonline_files(settings, mocker, mocked_celery, mock_block
     )
 
     with pytest.raises(mocked_celery.replace_exception_class):
-        tasks.import_all_mitxonline_files.delay(chunk_size=3, overwrite=True)
+        tasks.import_all_mitxonline_files.delay(
+            chunk_size=3, overwrite=True, learning_resource_ids=None
+        )
     get_content_tasks_mock.assert_called_once_with(
-        PlatformType.mitxonline.name, chunk_size=3, overwrite=True
+        PlatformType.mitxonline.name,
+        chunk_size=3,
+        overwrite=True,
+        learning_resource_ids=None,
     )
 
 
@@ -177,9 +185,9 @@ def test_import_all_xpro_files(settings, mocker, mocked_celery, mock_blocklist):
         "learning_resources.tasks.get_content_tasks", autospec=True
     )
     with pytest.raises(mocked_celery.replace_exception_class):
-        tasks.import_all_xpro_files.delay(chunk_size=3)
+        tasks.import_all_xpro_files.delay(chunk_size=3, learning_resource_ids=[1])
     get_content_tasks_mock.assert_called_once_with(
-        PlatformType.xpro.name, chunk_size=3, overwrite=False
+        PlatformType.xpro.name, chunk_size=3, overwrite=False, learning_resource_ids=[1]
     )
 
 
@@ -198,11 +206,19 @@ def test_import_all_oll_files(settings, mocker, mocked_celery, mock_blocklist):
         s3_prefix="open-learning-library/courses",
         override_base_prefix=True,
         overwrite=False,
+        learning_resource_ids=None,
     )
 
 
 @mock_aws
-def test_get_content_tasks(settings, mocker, mocked_celery, mock_xpro_learning_bucket):
+@pytest.mark.parametrize("with_learning_resource_ids", [True, False])
+def test_get_content_tasks(
+    settings,
+    mocker,
+    mocked_celery,
+    mock_xpro_learning_bucket,
+    with_learning_resource_ids,
+):
     """Test that get_content_tasks calls get_content_files with the correct args"""
     mock_get_content_files = mocker.patch(
         "learning_resources.tasks.get_content_files.si"
@@ -216,9 +232,23 @@ def test_get_content_tasks(settings, mocker, mocked_celery, mock_xpro_learning_b
     settings.LEARNING_COURSE_ITERATOR_CHUNK_SIZE = 2
     etl_source = ETLSource.xpro.name
     platform = PlatformType.xpro.name
-    factories.CourseFactory.create_batch(3, etl_source=etl_source, platform=platform)
+    courses = factories.CourseFactory.create_batch(
+        3, etl_source=etl_source, platform=platform
+    )
+    if with_learning_resource_ids:
+        learning_resource_ids = [
+            courses[0].learning_resource_id,
+            courses[1].learning_resource_id,
+        ]
+    else:
+        learning_resource_ids = None
     s3_prefix = "course-prefix"
-    tasks.get_content_tasks(etl_source, s3_prefix=s3_prefix, overwrite=True)
+    tasks.get_content_tasks(
+        etl_source,
+        s3_prefix=s3_prefix,
+        overwrite=True,
+        learning_resource_ids=learning_resource_ids,
+    )
     assert mocked_celery.group.call_count == 1
     assert (
         models.LearningResource.objects.filter(
@@ -230,10 +260,20 @@ def test_get_content_tasks(settings, mocker, mocked_celery, mock_xpro_learning_b
         .order_by("id")
         .values_list("id", flat=True)
     ).count() == 3
-    assert mock_get_content_files.call_count == 2
-    mock_get_content_files.assert_any_call(
-        ANY, etl_source, ["foo.tar.gz"], s3_prefix=s3_prefix, overwrite=True
-    )
+    if with_learning_resource_ids:
+        assert mock_get_content_files.call_count == 1
+        mock_get_content_files.assert_any_call(
+            [learning_resource_ids[0], learning_resource_ids[1]],
+            etl_source,
+            ["foo.tar.gz"],
+            s3_prefix=s3_prefix,
+            overwrite=True,
+        )
+    else:
+        assert mock_get_content_files.call_count == 2
+        mock_get_content_files.assert_any_call(
+            ANY, etl_source, ["foo.tar.gz"], s3_prefix=s3_prefix, overwrite=True
+        )
 
 
 def test_get_content_files(mocker, mock_xpro_learning_bucket):