Convert importability_annotation.py and the content/upgrade.py rollups

The temporary annotation in get_channel_annotation_stats rolls back through
transaction.set_rollback rather than a bridge transaction, and reads land as ORM
values_list, so coerce_key is no longer needed to turn PostgreSQL uuids back
into hex. The two historical rollup hooks differ only in the column they roll
up, so they share one helper.

sqlalchemy.exc.DatabaseError stays: import_external_content_dbs catches it
alongside sqlite3.DatabaseError around import_channel_from_local_db, which still
goes through channel_import.py.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: rtibblesbot

kolibri/core/content/test/test_importability.py

@@ -1,26 +1,22 @@
 from django.core.management import call_command
 from django.test import TransactionTestCase
-from mock import patch
+from le_utils.constants import content_kinds
 
+from kolibri.core.content.models import ContentNode
 from kolibri.core.content.models import File
 from kolibri.core.content.models import LocalFile
+from kolibri.core.content.utils.channels import CHANNEL_UPDATE_STATS_CACHE_KEY
+from kolibri.core.content.utils.content_types_tools import renderable_files_presets
 from kolibri.core.content.utils.importability_annotation import (
     get_channel_annotation_stats,
 )
-
-from .sqlalchemytesting import django_connection_engine
-
-
-def get_engine(connection_string):
-    return django_connection_engine()
-
+from kolibri.core.utils.cache import process_cache
 
 test_channel_id = "6199dde695db4ee4ab392222d5af1e5c"
 file_id_1 = "6bdfea4a01830fdd4a585181c0b8068c"
 file_id_2 = "e00699f859624e0f875ac6fe1e13d648"
 
 
-@patch("kolibri.core.content.utils.sqlalchemybridge.get_engine", new=get_engine)
 class ImportabilityStats(TransactionTestCase):
     fixtures = ["content_test.json"]
 
@@ -37,6 +33,69 @@ def test_all_files_available_no_files_remote(self):
         stats = get_channel_annotation_stats(test_channel_id, checksums)
         self.assertEqual(len(stats), 4)
 
+    def test_root_stats_count_every_importable_resource(self):
+        File.objects.update(supplementary=False)
+        checksums = list(LocalFile.objects.all().values_list("id", flat=True))
+        root = ContentNode.objects.get(channel_id=test_channel_id, level=0)
+        # Counted off the tree, not by a second rollup query, so that a rollup that
+        # summed over all children or lost its correlation cannot agree with it.
+        expected = (
+            root.get_descendants()
+            .exclude(kind=content_kinds.TOPIC)
+            .filter(
+                files__supplementary=False, files__preset__in=renderable_files_presets
+            )
+            .distinct()
+            .count()
+        )
+        stats = get_channel_annotation_stats(test_channel_id, checksums)
+        self.assertEqual(expected, stats[root.id]["total_resources"])
+
+    def test_the_projection_is_rolled_back(self):
+        File.objects.update(supplementary=False)
+        checksums = list(LocalFile.objects.all().values_list("id", flat=True))
+        before = self._annotation_columns()
+        get_channel_annotation_stats(test_channel_id, checksums)
+        self.assertEqual(before, self._annotation_columns())
+
+    def test_new_resources_are_flagged_on_the_root(self):
+        File.objects.update(supplementary=False)
+        checksums = list(LocalFile.objects.all().values_list("id", flat=True))
+        root = ContentNode.objects.get(channel_id=test_channel_id, level=0)
+        # A leaf the projection will call available, so that it reaches the stats dict.
+        new_ids = list(
+            root.get_descendants()
+            .exclude(kind=content_kinds.TOPIC)
+            .filter(
+                files__supplementary=False, files__preset__in=renderable_files_presets
+            )
+            .distinct()
+            .values_list("id", flat=True)[:1]
+        )
+        key = CHANNEL_UPDATE_STATS_CACHE_KEY.format(test_channel_id)
+        process_cache.set(
+            key, {"new_resource_ids": new_ids, "updated_resource_ids": new_ids}, None
+        )
+        self.addCleanup(process_cache.delete, key)
+        before = self._annotation_columns()
+        stats = get_channel_annotation_stats(test_channel_id, checksums)
+        self.assertTrue(stats[root.id]["new_resource"])
+        self.assertEqual(1, stats[root.id]["num_new_resources"])
+        self.assertTrue(stats[new_ids[0]]["updated_resource"])
+        # This block sets every node in the channel unavailable before measuring.
+        self.assertEqual(before, self._annotation_columns())
+
+    def _annotation_columns(self):
+        return list(
+            ContentNode.objects.order_by("id").values_list(
+                "id",
+                "available",
+                "coach_content",
+                "num_coach_contents",
+                "on_device_resources",
+            )
+        )
+
     def tearDown(self):
         call_command("flush", interactive=False)
         super().tearDown()

kolibri/core/content/test/test_upgrade.py

@@ -18,18 +18,12 @@
 from kolibri.core.content.upgrade import fix_multiple_trees_with_tree_id1
 from kolibri.core.content.upgrade import migrate_file_size_to_bigint
 from kolibri.core.content.upgrade import update_num_coach_contents
+from kolibri.core.content.upgrade import update_on_device_resources
 from kolibri.core.content.utils.content_types_tools import renderable_preset_bits
 from kolibri.core.content.utils.file_size_migration import drop_legacy_file_size_column
 from kolibri.core.content.utils.file_size_migration import localfile_columns
 from kolibri.core.content.utils.upgrade import diff_stats
 
-from .sqlalchemytesting import django_connection_engine
-
-
-def get_engine(connection_string):
-    return django_connection_engine()
-
-
 test_channel_id = "6199dde695db4ee4ab392222d5af1e5c"
 
 
@@ -186,7 +180,6 @@ def test_two_extra_channels_one_contentdb_exists(self, path_mock, import_mock):
         import_mock.assert_called_with(root_node_2.channel_id)
 
 
-@patch("kolibri.core.content.utils.sqlalchemybridge.get_engine", new=get_engine)
 class UpdateNumCoachContents(TransactionTestCase):
     fixtures = ["content_test.json"]
 
@@ -304,6 +297,37 @@ def tearDown(self):
         super().tearDown()
 
 
+class UpdateOnDeviceResources(TransactionTestCase):
+    fixtures = ["content_test.json"]
+
+    def setUp(self):
+        super().setUp()
+        ContentNode.objects.all().update(available=True, on_device_resources=0)
+        self.root_node = ContentNode.objects.get(parent__isnull=True)
+        self.leaves = (
+            self.root_node.get_descendants().exclude(kind=content_kinds.TOPIC).count()
+        )
+
+    def test_root_counts_its_available_leaf_descendants(self):
+        update_on_device_resources()
+        self.root_node.refresh_from_db()
+        self.assertEqual(self.leaves, self.root_node.on_device_resources)
+
+    def test_an_unavailable_leaf_is_not_counted(self):
+        leaf = (
+            self.root_node.get_descendants().exclude(kind=content_kinds.TOPIC).first()
+        )
+        leaf.available = False
+        leaf.save()
+        update_on_device_resources()
+        self.root_node.refresh_from_db()
+        self.assertEqual(self.leaves - 1, self.root_node.on_device_resources)
+
+    def tearDown(self):
+        call_command("flush", interactive=False)
+        super().tearDown()
+
+
 class FileIncludedPresetsAnnotationTestCase(TestCase):
     """
     Verifies the upgrade task backfills File.included_presets for rows that

kolibri/core/content/upgrade.py

@@ -7,18 +7,18 @@
 import sqlite3
 
 from django.db import connection
+from django.db import transaction
+from django.db.models import IntegerField
+from django.db.models import Sum
+from django.db.models.functions import Cast
 from le_utils.constants import content_kinds
 from le_utils.constants import library as library_constants
-from sqlalchemy import and_
-from sqlalchemy import cast
-from sqlalchemy import exists
-from sqlalchemy import func
-from sqlalchemy import Integer
-from sqlalchemy import select
+
+# import_channel_from_local_db still goes through channel_import.py, which raises
+# SQLAlchemy errors.
 from sqlalchemy.exc import DatabaseError
 
 from kolibri.core.auth.models import FacilityDataset
-from kolibri.core.content.apps import KolibriContentConfig
 from kolibri.core.content.constants.kind_to_learningactivity import kind_activity_map
 from kolibri.core.content.kolibri_plugin import synchronize_content_requests
 from kolibri.core.content.models import ChannelMetadata
@@ -27,9 +27,11 @@
 from kolibri.core.content.models import LocalFile
 from kolibri.core.content.tasks import backfill_content_request_priority
 from kolibri.core.content.tasks import enqueue_automatic_resource_import_if_needed
+from kolibri.core.content.utils.annotation import available_children_rollup
 from kolibri.core.content.utils.annotation import calculate_included_languages
 from kolibri.core.content.utils.annotation import calculate_ordered_categories
 from kolibri.core.content.utils.annotation import calculate_ordered_grade_levels
+from kolibri.core.content.utils.annotation import has_available_children
 from kolibri.core.content.utils.annotation import set_channel_ancestors
 from kolibri.core.content.utils.annotation import set_content_visibility_from_disk
 from kolibri.core.content.utils.channel_import import FutureSchemaError
@@ -44,7 +46,6 @@
 from kolibri.core.content.utils.search import annotate_label_bitmasks
 from kolibri.core.content.utils.search import annotate_modality
 from kolibri.core.content.utils.search import get_all_contentnode_label_metadata
-from kolibri.core.content.utils.sqlalchemybridge import Bridge
 from kolibri.core.content.utils.tree import get_channel_node_depth
 from kolibri.core.device.models import ContentCacheKey
 from kolibri.core.upgrade import version_upgrade
@@ -145,6 +146,27 @@ def fix_multiple_trees_with_tree_id1():
             )
 
 
+def _rollup_over_available_children(field, leaf_value):
+    """
+    Set field on every leaf from leaf_value, then sum it up the tree.
+    """
+    with transaction.atomic():
+        ContentNode.objects.exclude(kind=content_kinds.TOPIC).update(
+            **{field: leaf_value}
+        )
+
+        for channel_id in ChannelMetadata.objects.all().values_list("id", flat=True):
+            # Go from the deepest level to the shallowest
+            for level in range(get_channel_node_depth(channel_id), 0, -1):
+                ContentNode.objects.filter(
+                    level=level - 1, channel_id=channel_id, kind=content_kinds.TOPIC
+                ).filter(
+                    # A sum over no available children is NULL, so leave those
+                    # topics at the value they already carry.
+                    has_available_children()
+                ).update(**{field: available_children_rollup(Sum(field))})
+
+
 # This was introduced in 0.12.4, so only annotate
 # when upgrading from versions prior to this.
 @version_upgrade(old_version="<0.12.4")
@@ -153,73 +175,12 @@ def update_num_coach_contents():
     Function to set num_coach_content on all topic trees to account for
     those that were imported before annotations were performed
     """
-    bridge = Bridge(app_name=KolibriContentConfig.label)
-
-    ContentNodeTable = bridge.get_table(ContentNode)
-
-    connection = bridge.get_connection()
-
-    child = ContentNodeTable.alias()
-
     logger.info("Updating num_coach_content on existing channels")
 
-    # start a transaction
-
-    trans = connection.begin()
-
-    # Update all leaf ContentNodes to have num_coach_content to 1 or 0
-    connection.execute(
-        ContentNodeTable.update()
-        .where(
-            # That are not topics
-            ContentNodeTable.c.kind != content_kinds.TOPIC
-        )
-        .values(num_coach_contents=cast(ContentNodeTable.c.coach_content, Integer()))
-    )
-
-    # Expression to capture all available child nodes of a contentnode
-    available_nodes = select(child.c.available).where(
-        and_(
-            child.c.available == True,  # noqa
-            ContentNodeTable.c.id == child.c.parent_id,
-        )
+    _rollup_over_available_children(
+        "num_coach_contents", Cast("coach_content", IntegerField())
     )
 
-    # Expression that sums the total number of coach contents for each child node
-    # of a contentnode
-    coach_content_num = select(func.sum(child.c.num_coach_contents)).where(
-        and_(
-            child.c.available == True,  # noqa
-            ContentNodeTable.c.id == child.c.parent_id,
-        )
-    )
-
-    for channel_id in ChannelMetadata.objects.all().values_list("id", flat=True):
-        node_depth = get_channel_node_depth(channel_id)
-
-        # Go from the deepest level to the shallowest
-        for level in range(node_depth, 0, -1):
-            # Only modify topic availability here
-            connection.execute(
-                ContentNodeTable.update()
-                .where(
-                    and_(
-                        ContentNodeTable.c.level == level - 1,
-                        ContentNodeTable.c.channel_id == channel_id,
-                        ContentNodeTable.c.kind == content_kinds.TOPIC,
-                    )
-                )
-                # Because we have set availability to False on all topics as a starting point
-                # we only need to make updates to topics with available children.
-                .where(exists(available_nodes))
-                .values(num_coach_contents=coach_content_num.scalar_subquery())
-            )
-
-    # commit the transaction
-    trans.commit()
-
-    bridge.end()
-
 
 # This was introduced in 0.13.0, so only annotate
 # when upgrading from versions prior to this.
@@ -229,73 +190,12 @@ def update_on_device_resources():
     Function to set on_device_resource on all topic trees to account for
     those that were imported before annotations were performed
     """
-    bridge = Bridge(app_name=KolibriContentConfig.label)
-
-    ContentNodeTable = bridge.get_table(ContentNode)
-
-    connection = bridge.get_connection()
-
-    child = ContentNodeTable.alias()
-
     logger.info("Updating on_device_resource on existing channels")
 
-    # start a transaction
-
-    trans = connection.begin()
-
-    # Update all leaf ContentNodes to have on_device_resource to 1 or 0
-    connection.execute(
-        ContentNodeTable.update()
-        .where(
-            # That are not topics
-            ContentNodeTable.c.kind != content_kinds.TOPIC
-        )
-        .values(on_device_resources=cast(ContentNodeTable.c.available, Integer()))
-    )
-
-    # Expression to capture all available child nodes of a contentnode
-    available_nodes = select(child.c.available).where(
-        and_(
-            child.c.available == True,  # noqa
-            ContentNodeTable.c.id == child.c.parent_id,
-        )
-    )
-
-    # Expression that sums the total number of coach contents for each child node
-    # of a contentnode
-    on_device_num = select(func.sum(child.c.on_device_resources)).where(
-        and_(
-            child.c.available == True,  # noqa
-            ContentNodeTable.c.id == child.c.parent_id,
-        )
+    _rollup_over_available_children(
+        "on_device_resources", Cast("available", IntegerField())
     )
 
-    for channel_id in ChannelMetadata.objects.all().values_list("id", flat=True):
-        node_depth = get_channel_node_depth(channel_id)
-
-        # Go from the deepest level to the shallowest
-        for level in range(node_depth, 0, -1):
-            # Only modify topic availability here
-            connection.execute(
-                ContentNodeTable.update()
-                .where(
-                    and_(
-                        ContentNodeTable.c.level == level - 1,
-                        ContentNodeTable.c.channel_id == channel_id,
-                        ContentNodeTable.c.kind == content_kinds.TOPIC,
-                    )
-                )
-                # Because we have set availability to False on all topics as a starting point
-                # we only need to make updates to topics with available children.
-                .where(exists(available_nodes))
-                .values(on_device_resources=on_device_num)
-            )
-
-    # commit the transaction
-    trans.commit()
-
-    bridge.end()
-
 
 # This was introduced in 0.15.0, so only annotate
 # when upgrading from versions prior to this.