DRY things out a bit and share as much code between MongoModuleStore and DraftMongoModuleStore

common/lib/xmodule/xmodule/modulestore/draft.py

@@ -191,45 +191,18 @@ class DraftModuleStore(ModuleStoreBase):
         super(DraftModuleStore, self).clone_item(location, as_draft(location))
         super(DraftModuleStore, self).delete_item(location)
 
-    def _cache_children(self, items, depth=0):
-        """
-        Returns a dictionary mapping Location -> item data, populated with json data
-        for all descendents of items up to the specified depth.
-        (0 = no descendents, 1 = children, 2 = grandchildren, etc)
-        If depth is None, will load all the children.
-        This will make a number of queries that is linear in the depth.
-        """
-
-        data = {}
-        to_process = list(items)
-        while to_process and depth is None or depth >= 0:
-            children = []
-            for item in to_process:
-                self._clean_item_data(item)
-                children.extend(item.get('definition', {}).get('children', []))
-                data[Location(item['location'])] = item
-
-            if depth == 0:
-                break;
-
-            # Load all children by id. See
-            # http://www.mongodb.org/display/DOCS/Advanced+Queries#AdvancedQueries-%24or
-            # for or-query syntax
-            to_process = []
-            if children:
+    def _query_children_for_cache_children(self, items):
         # first get non-draft in a round-trip
-                query = {
-                    '_id': {'$in': [namedtuple_to_son(Location(child)) for child in children]}
-                }
-                to_process_non_drafts = list(self.collection.find(query))
+        queried_children = []
+        to_process_non_drafts = super(DraftModuleStore, self)._query_children_for_cache_children(items)
 
         to_process_dict = {}
         for non_draft in to_process_non_drafts:
             to_process_dict[Location(non_draft["_id"])] = non_draft
 
-                # now query all draft content in a round-trip
+        # now query all draft content in another round-trip
         query = {
-                    '_id': {'$in': [namedtuple_to_son(as_draft(Location(child))) for child in children]}
+            '_id': {'$in': [namedtuple_to_son(as_draft(Location(item))) for item in items]}
         }
         to_process_drafts = list(self.collection.find(query))
 
@@ -247,10 +220,6 @@ class DraftModuleStore(ModuleStoreBase):
 
         # convert the dict - which is used for look ups - back into a list
         for key, value in to_process_dict.iteritems():
-                    to_process.append(value)
-
-            # If depth is None, then we just recurse until we hit all the descendents
-            if depth is not None:
-                depth -= 1
+            queried_children.append(value)  
 
-        return data
\ No newline at end of file
+        return queried_children

common/lib/xmodule/xmodule/modulestore/mongo.py

@@ -363,6 +363,13 @@ class MongoModuleStore(ModuleStoreBase):
         item['location'] = item['_id']
         del item['_id']
 
+    def _query_children_for_cache_children(self, items):
+        # first get non-draft in a round-trip
+        query = {
+            '_id': {'$in': [namedtuple_to_son(Location(item)) for item in items]}
+        }
+        return list(self.collection.find(query))
+
     def _cache_children(self, items, depth=0):
         """
         Returns a dictionary mapping Location -> item data, populated with json data
@@ -382,18 +389,14 @@ class MongoModuleStore(ModuleStoreBase):
                 data[Location(item['location'])] = item
 
             if depth == 0:
-                break
+                break;
 
             # Load all children by id. See
             # http://www.mongodb.org/display/DOCS/Advanced+Queries#AdvancedQueries-%24or
             # for or-query syntax
-            if children and depth > 0:
-                query = {
-                    '_id': {'$in': [namedtuple_to_son(Location(child)) for child in children]}
-                }
-                to_process = list(self.collection.find(query))
-            else:
-                break
+            to_process = []
+            if children:
+                to_process = self._query_children_for_cache_children(children)
 
             # If depth is None, then we just recurse until we hit all the descendents
             if depth is not None: