perf: select related objects in container queries (#333)

Author: pomegranited

openedx_learning/__init__.py

@@ -2,4 +2,4 @@
 Open edX Learning ("Learning Core").
 """
 
-__version__ = "0.27.0"
+__version__ = "0.27.1"

openedx_learning/apps/authoring/publishing/api.py

@@ -1283,6 +1283,7 @@ def get_entities_in_container(
     container: Container,
     *,
     published: bool,
+    select_related_version: str | None = None,
 ) -> list[ContainerEntityListEntry]:
     """
     [ 🛑 UNSTABLE ]
@@ -1293,14 +1294,35 @@ def get_entities_in_container(
         container: The Container, e.g. returned by `get_container()`
         published: `True` if we want the published version of the container, or
             `False` for the draft version.
+        select_related_version: An optional optimization; specify a relationship
+        on ContainerVersion, like `componentversion` or `containerversion__x`
+        to preload via select_related.
     """
     assert isinstance(container, Container)
-    container_version = container.versioning.published if published else container.versioning.draft
+    if published:
+        # Very minor optimization: reload the container with related 1:1 entities
+        container = Container.objects.select_related(
+            "publishable_entity__published__version__containerversion__entity_list").get(pk=container.pk)
+        container_version = container.versioning.published
+        select_related = ["entity__published__version"]
+        if select_related_version:
+            select_related.append(f"entity__published__version__{select_related_version}")
+    else:
+        # Very minor optimization: reload the container with related 1:1 entities
+        container = Container.objects.select_related(
+            "publishable_entity__draft__version__containerversion__entity_list").get(pk=container.pk)
+        container_version = container.versioning.draft
+        select_related = ["entity__draft__version"]
+        if select_related_version:
+            select_related.append(f"entity__draft__version__{select_related_version}")
     if container_version is None:
         raise ContainerVersion.DoesNotExist  # This container has not been published yet, or has been deleted.
     assert isinstance(container_version, ContainerVersion)
-    entity_list = []
-    for row in container_version.entity_list.entitylistrow_set.order_by("order_num"):
+    entity_list: list[ContainerEntityListEntry] = []
+    for row in container_version.entity_list.entitylistrow_set.select_related(
+        "entity_version",
+        *select_related,
+    ).order_by("order_num"):
         entity_version = row.entity_version  # This will be set if pinned
         if not entity_version:  # If this entity is "unpinned", use the latest published/draft version:
             entity_version = row.entity.published.version if published else row.entity.draft.version
@@ -1393,7 +1415,10 @@ def get_containers_with_entity(
         qs = Container.objects.filter(
             publishable_entity__draft__version__containerversion__entity_list__entitylistrow__entity_id=publishable_entity_pk,  # pylint: disable=line-too-long # noqa: E501
         )
-    return qs.order_by("pk").distinct()  # Ordering is mostly for consistent test cases.
+    return qs.select_related(
+        "publishable_entity__draft__version__containerversion",
+        "publishable_entity__published__version__containerversion",
+    ).order_by("pk").distinct()  # Ordering is mostly for consistent test cases.
 
 
 def get_container_children_count(

openedx_learning/apps/authoring/sections/api.py

@@ -257,7 +257,12 @@ def get_subsections_in_section(
     """
     assert isinstance(section, Section)
     subsections = []
-    for entry in publishing_api.get_entities_in_container(section, published=published):
+    entries = publishing_api.get_entities_in_container(
+        section,
+        published=published,
+        select_related_version="containerversion__subsectionversion",
+    )
+    for entry in entries:
         # Convert from generic PublishableEntityVersion to SubsectionVersion:
         subsection_version = entry.entity_version.containerversion.subsectionversion
         assert isinstance(subsection_version, SubsectionVersion)

openedx_learning/apps/authoring/subsections/api.py

@@ -256,7 +256,12 @@ def get_units_in_subsection(
     """
     assert isinstance(subsection, Subsection)
     units = []
-    for entry in publishing_api.get_entities_in_container(subsection, published=published):
+    entries = publishing_api.get_entities_in_container(
+        subsection,
+        published=published,
+        select_related_version="containerversion__unitversion",
+    )
+    for entry in entries:
         # Convert from generic PublishableEntityVersion to UnitVersion:
         unit_version = entry.entity_version.containerversion.unitversion
         assert isinstance(unit_version, UnitVersion)

openedx_learning/apps/authoring/units/api.py

@@ -257,7 +257,12 @@ def get_components_in_unit(
     """
     assert isinstance(unit, Unit)
     components = []
-    for entry in publishing_api.get_entities_in_container(unit, published=published):
+    entries = publishing_api.get_entities_in_container(
+        unit,
+        published=published,
+        select_related_version="componentversion",
+    )
+    for entry in entries:
         # Convert from generic PublishableEntityVersion to ComponentVersion:
         component_version = entry.entity_version.componentversion
         assert isinstance(component_version, ComponentVersion)

tests/openedx_learning/apps/authoring/sections/test_api.py

@@ -993,6 +993,32 @@ def test_sections_containing(self):
             ]
         assert result2 == [section4_unpinned]
 
+    def test_get_subsections_in_section_queries(self):
+        """
+        Test the query count of get_subsections_in_section()
+        This also tests the generic method get_entities_in_container()
+        """
+        section = self.create_section_with_subsections([
+            self.subsection_1,
+            self.subsection_2,
+            self.subsection_2_v1,
+        ])
+        with self.assertNumQueries(4):
+            result = authoring_api.get_subsections_in_section(section, published=False)
+        assert result == [
+            Entry(self.subsection_1.versioning.draft),
+            Entry(self.subsection_2.versioning.draft),
+            Entry(self.subsection_2.versioning.draft, pinned=True),
+        ]
+        authoring_api.publish_all_drafts(self.learning_package.id)
+        with self.assertNumQueries(4):
+            result = authoring_api.get_subsections_in_section(section, published=True)
+        assert result == [
+            Entry(self.subsection_1.versioning.draft),
+            Entry(self.subsection_2.versioning.draft),
+            Entry(self.subsection_2.versioning.draft, pinned=True),
+        ]
+
     def test_add_remove_container_children(self):
         """
         Test adding and removing children subsections from sections.

tests/openedx_learning/apps/authoring/subsections/test_api.py

@@ -1024,6 +1024,32 @@ def test_subsections_containing(self):
             ]
         assert result2 == [subsection4_unpinned]
 
+    def test_get_units_in_subsection_queries(self):
+        """
+        Test the query count of get_units_in_subsection()
+        This also tests the generic method get_entities_in_container()
+        """
+        subsection = self.create_subsection_with_units([
+            self.unit_1,
+            self.unit_2,
+            self.unit_2_v1,
+        ])
+        with self.assertNumQueries(4):
+            result = authoring_api.get_units_in_subsection(subsection, published=False)
+        assert result == [
+            Entry(self.unit_1.versioning.draft),
+            Entry(self.unit_2.versioning.draft),
+            Entry(self.unit_2.versioning.draft, pinned=True),
+        ]
+        authoring_api.publish_all_drafts(self.learning_package.id)
+        with self.assertNumQueries(4):
+            result = authoring_api.get_units_in_subsection(subsection, published=True)
+        assert result == [
+            Entry(self.unit_1.versioning.draft),
+            Entry(self.unit_2.versioning.draft),
+            Entry(self.unit_2.versioning.draft, pinned=True),
+        ]
+
     def test_add_remove_container_children(self):
         """
         Test adding and removing children units from subsections.

tests/openedx_learning/apps/authoring/units/test_api.py

@@ -1012,6 +1012,32 @@ def test_units_containing(self):
             ]
         assert result2 == [unit4_unpinned, unit7_several]
 
+    def test_get_components_in_unit_queries(self):
+        """
+        Test the query count of get_components_in_unit()
+        This also tests the generic method get_entities_in_container()
+        """
+        unit = self.create_unit_with_components([
+            self.component_1,
+            self.component_2,
+            self.component_2_v1,
+        ])
+        with self.assertNumQueries(3):
+            result = authoring_api.get_components_in_unit(unit, published=False)
+        assert result == [
+            Entry(self.component_1.versioning.draft),
+            Entry(self.component_2.versioning.draft),
+            Entry(self.component_2.versioning.draft, pinned=True),
+        ]
+        authoring_api.publish_all_drafts(self.learning_package.id)
+        with self.assertNumQueries(3):
+            result = authoring_api.get_components_in_unit(unit, published=True)
+        assert result == [
+            Entry(self.component_1.versioning.draft),
+            Entry(self.component_2.versioning.draft),
+            Entry(self.component_2.versioning.draft, pinned=True),
+        ]
+
     def test_add_remove_container_children(self):
         """
         Test adding and removing children components from containers.