From e8fecd17c4ed36284cd22f4a1fcbed2ccd20c35c Mon Sep 17 00:00:00 2001
From: Braden MacDonald <braden@opencraft.com>
Date: Fri, 14 Feb 2025 11:07:40 -0800
Subject: [PATCH] Optimize contains_unpublished_changes to check a unit using
 only 3 queries

---
 .../apps/authoring/containers/api.py          | 17 +++++++-
 .../apps/authoring/units/test_api.py          | 39 ++++++++++++++++++-
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/openedx_learning/apps/authoring/containers/api.py b/openedx_learning/apps/authoring/containers/api.py
index d67126cd..cc9b60ff 100644
--- a/openedx_learning/apps/authoring/containers/api.py
+++ b/openedx_learning/apps/authoring/containers/api.py
@@ -551,7 +551,16 @@ def contains_unpublished_changes(
     itself has unpublished changes, not if its contents do.
     """
     if isinstance(container, ContainerEntityMixin):
-        container = container.container_entity
+        # The query below pre-loads the data we need but is otherwise the same thing as:
+        #     container = container.container_entity
+        container = ContainerEntity.objects.select_related(
+            "publishable_entity",
+            "publishable_entity__draft",
+            "publishable_entity__draft__version",
+            "publishable_entity__draft__version__containerentityversion__defined_list",
+        ).get(pk=container.container_entity)
+    else:
+        pass # TODO: select_related if we're given a raw ContainerEntity rather than a ContainerEntityMixin like Unit?
     assert isinstance(container, ContainerEntity)
 
     if container.versioning.has_unpublished_changes:
@@ -563,7 +572,11 @@ def contains_unpublished_changes(
     # TODO: This is a naive inefficient implementation but hopefully correct.
     # Once we know it's correct and have a good test suite, then we can optimize.
     # We will likely change to a tracking-based approach rather than a "scan for changes" based approach.
-    for row in defined_list.entitylistrow_set.filter(draft_version=None):
+    for row in defined_list.entitylistrow_set.filter(draft_version=None).select_related(
+        "entity__containerentity",
+        "entity__draft__version",
+        "entity__published__version",
+    ):
         try:
             child_container = row.entity.containerentity
         except PublishableEntity.containerentity.RelatedObjectDoesNotExist:
diff --git a/tests/openedx_learning/apps/authoring/units/test_api.py b/tests/openedx_learning/apps/authoring/units/test_api.py
index 22fa1807..abcfe5ed 100644
--- a/tests/openedx_learning/apps/authoring/units/test_api.py
+++ b/tests/openedx_learning/apps/authoring/units/test_api.py
@@ -205,7 +205,44 @@ def test_modify_component_after_publish(self):
         assert authoring_api.contains_unpublished_changes(unit) == False  # No longer contains unpublished changes
 
 
-    # Test query count of contains_unpublished_changes()
+    def test_query_count_of_contains_unpublished_changes(self):
+        """
+        Checking for unpublished changes in a unit should require a fixed number
+        of queries, not get more expensive as the unit gets larger.
+        """
+        unit, unit_version = authoring_api.create_unit_and_version(
+            learning_package_id=self.learning_package.id,
+            key=f"unit:key",
+            title="Unit",
+            created=self.now,
+            created_by=None,
+        )
+        # Add 100 components (unpinned)
+        component_count = 100
+        publishable_entities_pks = []
+        for i in range(0, component_count):
+            component, _version = authoring_api.create_component_and_version(
+                self.learning_package.id,
+                component_type=self.problem_type,
+                local_key=f"Query Counting {i}",
+                title=f"Querying Counting Problem {i}",
+                created=self.now,
+            )
+            publishable_entities_pks.append(component.publishable_entity_id)
+        authoring_api.create_next_unit_version(
+            unit=unit,
+            title=unit_version.title,
+            publishable_entities_pks=publishable_entities_pks,
+            draft_version_pks=[None] * component_count,
+            published_version_pks=[None] * component_count,  # FIXME: why do we specify this?
+            created=self.now,
+        )
+        authoring_api.publish_all_drafts(self.learning_package.id)
+        unit.refresh_from_db()
+        with self.assertNumQueries(3):
+            assert authoring_api.contains_unpublished_changes(unit) == False
+
+    # Test that pinned components with changes don't show up as "contains unpublished changes"
     # Test that only components can be added to units
     # Test that components must be in the same learning package
     # Test that _version_pks=[] arguments must be related to publishable_entities_pks