Move query-chunking into StudentModule and related ORM-objects

lms/djangoapps/courseware/model_data.py

@@ -5,7 +5,6 @@ Classes to provide the LMS runtime data storage to XBlocks
 import json
 from abc import abstractmethod, ABCMeta
 from collections import defaultdict
-from itertools import chain
 from .models import (
     StudentModule,
     XModuleUserStateSummaryField,
@@ -36,29 +35,6 @@ class InvalidWriteError(Exception):
     """
 
 
-def chunks(items, chunk_size):
-    """
-    Yields the values from items in chunks of size chunk_size
-    """
-    items = list(items)
-    return (items[i:i + chunk_size] for i in xrange(0, len(items), chunk_size))
-
-
-def _chunked_query(model_class, chunk_field, items, chunk_size=500, **kwargs):
-    """
-    Queries model_class with `chunk_field` set to chunks of size `chunk_size`,
-    and all other parameters from `**kwargs`.
-
-    This works around a limitation in sqlite3 on the number of parameters
-    that can be put into a single query.
-    """
-    res = chain.from_iterable(
-        model_class.objects.filter(**dict([(chunk_field, chunk)] + kwargs.items()))
-        for chunk in chunks(items, chunk_size)
-    )
-    return res
-
-
 def _all_usage_keys(descriptors, aside_types):
     """
     Return a set of all usage_ids for the `descriptors` and for
@@ -362,8 +338,7 @@ class UserStateCache(object):
             xblocks (list of :class:`XBlock`): XBlocks to cache fields for.
             aside_types (list of str): Aside types to cache fields for.
         """
-        query = _chunked_query(
-            StudentModule,
+        query = StudentModule.objects.chunked_filter(
             'module_state_key__in',
             _all_usage_keys(xblocks, aside_types),
             course_id=self.course_id,
@@ -567,8 +542,7 @@ class UserStateSummaryCache(DjangoOrmFieldCache):
             aside_types (list of str): Asides to load field for (which annotate the supplied
                 xblocks).
         """
-        return _chunked_query(
-            XModuleUserStateSummaryField,
+        return XModuleUserStateSummaryField.objects.chunked_filter(
             'usage_id__in',
             _all_usage_keys(xblocks, aside_types),
             field_name__in=set(field.name for field in fields),
@@ -629,8 +603,7 @@ class PreferencesCache(DjangoOrmFieldCache):
             aside_types (list of str): Asides to load field for (which annotate the supplied
                 xblocks).
         """
-        return _chunked_query(
-            XModuleStudentPrefsField,
+        return XModuleStudentPrefsField.objects.chunked_filter(
             'module_type__in',
             _all_block_types(xblocks, aside_types),
             student=self.user.pk,

lms/djangoapps/courseware/models.py

@@ -13,6 +13,7 @@ ASSUMPTIONS: modules have unique IDs, even across different module_types
 
 """
 import logging
+import itertools
 
 from django.contrib.auth.models import User
 from django.conf import settings
@@ -29,10 +30,49 @@ from xmodule_django.models import CourseKeyField, LocationKeyField, BlockTypeKey
 log = logging.getLogger("edx.courseware")
 
 
+def chunks(items, chunk_size):
+    """
+    Yields the values from items in chunks of size chunk_size
+    """
+    items = list(items)
+    return (items[i:i + chunk_size] for i in xrange(0, len(items), chunk_size))
+
+
+class ChunkingManager(models.Manager):
+    """
+    :class:`~Manager` that adds an additional method :meth:`chunked_filter` to provide
+    the ability to make select queries with specific chunk sizes.
+    """
+    def chunked_filter(self, chunk_field, items, **kwargs):
+        """
+        Queries model_class with `chunk_field` set to chunks of size `chunk_size`,
+        and all other parameters from `**kwargs`.
+
+        This works around a limitation in sqlite3 on the number of parameters
+        that can be put into a single query.
+
+        Arguments:
+            chunk_field (str): The name of the field to chunk the query on.
+            items: The values for of chunk_field to select. This is chunked into ``chunk_size``
+                chunks, and passed as the value for the ``chunk_field`` keyword argument to
+                :meth:`~Manager.filter`. This implies that ``chunk_field`` should be an
+                ``__in`` key.
+            chunk_size (int): The size of chunks to pass. Defaults to 500.
+        """
+        chunk_size = kwargs.pop('chunk_size', 500)
+        res = itertools.chain.from_iterable(
+            self.filter(**dict([(chunk_field, chunk)] + kwargs.items()))
+            for chunk in chunks(items, chunk_size)
+        )
+        return res
+
+
 class StudentModule(models.Model):
     """
     Keeps student state for a particular module in a particular course.
     """
+    objects = ChunkingManager()
+
     MODEL_TAGS = ['course_id', 'module_type']
 
     # For a homework problem, contains a JSON
@@ -142,6 +182,8 @@ class XBlockFieldBase(models.Model):
     """
     Base class for all XBlock field storage.
     """
+    objects = ChunkingManager()
+
     class Meta(object):  # pylint: disable=missing-docstring
         abstract = True