refactor: cleanup code and add query optimizations

Author: dwong2708

openedx_learning/apps/authoring/backup_restore/zipper.py

@@ -5,19 +5,21 @@
 import hashlib
 import zipfile
 from pathlib import Path
-from typing import Optional
+from typing import List, Optional
 
-from django.db.models import QuerySet
+from django.db.models import Prefetch, QuerySet
 from django.utils.text import slugify
 
-from openedx_learning.apps.authoring.backup_restore.toml import toml_learning_package, toml_publishable_entity
-from openedx_learning.apps.authoring.components.models import ComponentVersion, ComponentVersionContent
-from openedx_learning.apps.authoring.publishing import api as publishing_api
-from openedx_learning.apps.authoring.publishing.models import (
+from openedx_learning.api.authoring_models import (
+    ComponentVersion,
+    ComponentVersionContent,
+    Content,
     LearningPackage,
     PublishableEntity,
     PublishableEntityVersion,
 )
+from openedx_learning.apps.authoring.backup_restore.toml import toml_learning_package, toml_publishable_entity
+from openedx_learning.apps.authoring.publishing import api as publishing_api
 
 TOML_PACKAGE_NAME = "package.toml"
 
@@ -69,6 +71,49 @@ def create_folder(self, folder_path: Path, zip_file: zipfile.ZipFile) -> None:
             zip_file.writestr(zip_info, "")  # Add explicit empty directory entry
             self.folders_already_created.add(folder_path)
 
+    def get_publishable_entities(self) -> QuerySet[PublishableEntity]:
+        """
+        Retrieve the publishable entities associated with the learning package.
+        Prefetches related data for efficiency.
+        """
+        lp_id = self.learning_package.pk
+        publishable_entities: QuerySet[PublishableEntity] = publishing_api.get_publishable_entities(lp_id)
+        return (
+            publishable_entities
+            .select_related(
+                "container",
+                "component__component_type",
+                "draft__version__componentversion",
+                "published__version__componentversion",
+            )
+            .prefetch_related(
+                Prefetch(
+                    "draft__version__componentversion__componentversioncontent_set",
+                    queryset=ComponentVersionContent.objects.select_related("content"),
+                    to_attr="prefetched_contents",
+                ),
+                Prefetch(
+                    "published__version__componentversion__componentversioncontent_set",
+                    queryset=ComponentVersionContent.objects.select_related("content"),
+                    to_attr="prefetched_contents",
+                ),
+            )
+        )
+
+    def get_versions_to_write(self, entity: PublishableEntity):
+        """
+        Get the versions of a publishable entity that should be written to the zip file.
+        It retrieves both draft and published versions.
+        """
+        draft_version: Optional[PublishableEntityVersion] = publishing_api.get_draft_version(entity)
+        published_version: Optional[PublishableEntityVersion] = publishing_api.get_published_version(entity)
+
+        versions_to_write = [draft_version] if draft_version else []
+
+        if published_version and published_version != draft_version:
+            versions_to_write.append(published_version)
+        return versions_to_write
+
     def create_zip(self, path: str) -> None:
         """
         Creates a zip file containing the learning package data.
@@ -77,7 +122,6 @@ def create_zip(self, path: str) -> None:
         Raises:
             Exception: If the learning package cannot be found or if the zip creation fails.
         """
-        lp_id = self.learning_package.pk
 
         with zipfile.ZipFile(path, "w", compression=zipfile.ZIP_DEFLATED) as zipf:
             # Add the package.toml file
@@ -95,8 +139,7 @@ def create_zip(self, path: str) -> None:
             # ------ ENTITIES SERIALIZATION -------------
 
             # get the publishable entities
-            publishable_entities: QuerySet[PublishableEntity] = publishing_api.get_publishable_entities(lp_id)
-            publishable_entities = publishable_entities.select_related("container", "component__component_type")
+            publishable_entities: QuerySet[PublishableEntity] = self.get_publishable_entities()
 
             for entity in publishable_entities:
                 # entity: PublishableEntity = entity  # Type hint for clarity
@@ -155,13 +198,7 @@ def create_zip(self, path: str) -> None:
                     # Focusing on draft and published versions
 
                     # Get the draft and published versions
-                    draft_version: Optional[PublishableEntityVersion] = publishing_api.get_draft_version(entity)
-                    published_version: Optional[PublishableEntityVersion] = publishing_api.get_published_version(entity)
-
-                    versions_to_write = [draft_version] if draft_version else []
-
-                    if published_version and published_version != draft_version:
-                        versions_to_write.append(published_version)
+                    versions_to_write: List[PublishableEntityVersion] = self.get_versions_to_write(entity)
 
                     for version in versions_to_write:
                         # Create a folder for the version
@@ -174,24 +211,27 @@ def create_zip(self, path: str) -> None:
                         self.create_folder(static_folder, zipf)
 
                         # ------ COMPONENT STATIC CONTENT -------------
-                        # Get component version
                         component_version: ComponentVersion = version.componentversion
 
                         # Get content data associated with this version
-                        # content_list: QuerySet[Content] = component_version.contents.all()
-                        content_list: QuerySet[ComponentVersionContent] = component_version.componentversioncontent_set.all()  # pylint: disable=line-too-long  # noqa: E501
+                        # type: ignore[attr-defined]
+                        content_list: QuerySet[ComponentVersionContent] = component_version.prefetched_contents
 
                         for component_version_content in content_list:
-                            content = component_version_content.content
+                            content: Content = component_version_content.content
+
+                            # Important: The component_version_content.key contains implicitly
+                            # the file name and the file extension
+                            file_path = version_folder / component_version_content.key
 
                             if content.has_file and content.path:
-                                # Add the file to the static folder
-                                # file_path = static_folder / content.path
-                                file_path = static_folder / component_version_content.key
+                                # If has_file, we pull it from the file system
                                 with content.read_file() as f:
                                     file_data = f.read()
-                                    zipf.writestr(str(file_path), file_data)
                             elif not content.has_file and content.text:
-                                # Create file for the text file according to the mime_type attr
-                                text_file_path = static_folder / component_version_content.key
-                                zipf.writestr(str(text_file_path), content.text)
+                                # Otherwise, we use the text content as the file data
+                                file_data = content.text
+                            else:
+                                # If no file and no text, we skip this content
+                                continue
+                            zipf.writestr(str(file_path), file_data)

tests/openedx_learning/apps/authoring/backup_restore/test_backup.py

@@ -19,7 +19,7 @@
 
 class LpDumpCommandTestCase(TestCase):
     """
-    Test serving static assets (Content files, via Component lookup).
+    Test the lp_dump management command.
     """
 
     learning_package: LearningPackage
@@ -28,7 +28,7 @@ class LpDumpCommandTestCase(TestCase):
     @classmethod
     def setUpTestData(cls):
         """
-        Initialize our content data
+        Initialize data for the whole TestCase
         """
 
         # Create a user for the test
@@ -86,7 +86,7 @@ def setUpTestData(cls):
         api.create_component_version_content(
             new_problem_version.pk,
             new_txt_content.pk,
-            key="content/hello.txt",
+            key="hello.txt",
         )
 
         # Create a Draft component, one in each learning package
@@ -115,7 +115,7 @@ def setUpTestData(cls):
         api.create_component_version_content(
             new_html_version.pk,
             cls.html_asset_content.id,
-            key="content/hello.html",
+            key="static/hello.html",
         )
 
         components = api.get_publishable_entities(cls.learning_package)
@@ -154,8 +154,8 @@ def check_zip_file_structure(self, zip_path: Path):
                 "entities/xblock.v1/problem/my_published_example_386dce.toml",
 
                 # Entity static content files
-                "entities/xblock.v1/html/my_draft_example_af06e1/component_versions/v2/static/content/hello.html",
-                "entities/xblock.v1/problem/my_published_example_386dce/component_versions/v2/static/content/hello.txt",
+                "entities/xblock.v1/html/my_draft_example_af06e1/component_versions/v2/static/hello.html",
+                "entities/xblock.v1/problem/my_published_example_386dce/component_versions/v2/hello.txt",
             ]
 
             expected_paths = expected_directories + expected_files