Move context_managers to base class, overridden by mixed.

cms/djangoapps/contentstore/management/commands/clone_course.py

@@ -3,7 +3,6 @@ Script for cloning a course
 """
 from django.core.management.base import BaseCommand, CommandError
 from xmodule.modulestore.django import modulestore
-from xmodule.modulestore.mixed import store_bulk_write_operations_on_course
 from student.roles import CourseInstructorRole, CourseStaffRole
 from opaque_keys.edx.keys import CourseKey
 from opaque_keys import InvalidKeyError
@@ -38,7 +37,7 @@ class Command(BaseCommand):
 
         print("Cloning course {0} to {1}".format(source_course_id, dest_course_id))
 
-        with store_bulk_write_operations_on_course(mstore, dest_course_id):
+        with mstore.bulk_write_operations(dest_course_id):
             if mstore.clone_course(source_course_id, dest_course_id, None):
                 print("copying User permissions...")
                 # purposely avoids auth.add_user b/c it doesn't have a caller to authorize

cms/djangoapps/contentstore/management/commands/migrate_to_split.py

-# pylint: disable=W0212
+# pylint: disable=protected-access
 
 """
 Django management command to migrate a course from the old Mongo modulestore

cms/djangoapps/contentstore/tests/test_clone_course.py

@@ -29,7 +29,7 @@ class CloneCourseTest(CourseTestCase):
         self.assertCoursesEqual(mongo_course1_id, mongo_course2_id)
 
         # 3. clone course (mongo -> split)
-        with self.store.set_default_store(ModuleStoreEnum.Type.split):
+        with self.store.default_store(ModuleStoreEnum.Type.split):
             split_course3_id = CourseLocator(
                 org="edx3", course="split3", run="2013_Fall", branch=ModuleStoreEnum.BranchName.draft
             )

cms/djangoapps/contentstore/tests/test_contentstore.py

@@ -26,7 +26,6 @@ from xmodule.contentstore.django import contentstore, _CONTENTSTORE
 from xmodule.contentstore.utils import restore_asset_from_trashcan, empty_asset_trashcan
 from xmodule.exceptions import NotFoundError, InvalidVersionError
 from xmodule.modulestore import ModuleStoreEnum
-from xmodule.modulestore.mixed import store_branch_setting
 from xmodule.modulestore.exceptions import ItemNotFoundError
 from xmodule.modulestore.inheritance import own_metadata
 from opaque_keys.edx.keys import UsageKey
@@ -873,7 +872,7 @@ class ContentStoreToyCourseTest(ContentStoreTestCase):
         mongo_store.collection.find = wrapper.find
 
         # set the branch to 'publish' in order to prevent extra lookups of draft versions
-        with store_branch_setting(mongo_store, ModuleStoreEnum.Branch.published_only):
+        with mongo_store.branch_setting(ModuleStoreEnum.Branch.published_only):
             course = mongo_store.get_course(course_id, depth=2)
 
             # make sure we haven't done too many round trips to DB

cms/djangoapps/contentstore/utils.py

@@ -14,7 +14,6 @@ from xmodule.contentstore.content import StaticContent
 from xmodule.contentstore.django import contentstore
 from xmodule.modulestore import ModuleStoreEnum
 from xmodule.modulestore.django import modulestore
-from xmodule.modulestore.mixed import store_bulk_write_operations_on_course
 from xmodule.modulestore.exceptions import ItemNotFoundError
 from opaque_keys.edx.locations import SlashSeparatedCourseKey, Location
 from xmodule.modulestore.store_utilities import delete_course
@@ -37,7 +36,7 @@ def delete_course_and_groups(course_id, commit=False):
     module_store = modulestore()
     content_store = contentstore()
 
-    with store_bulk_write_operations_on_course(module_store, course_id):
+    with module_store.bulk_write_operations(course_id):
         if delete_course(module_store, content_store, course_id, commit):
 
             print 'removing User permissions from course....'

common/lib/xmodule/xmodule/modulestore/__init__.py

@@ -10,6 +10,7 @@ import datetime
 
 from collections import namedtuple, defaultdict
 import collections
+from contextlib import contextmanager
 
 from abc import ABCMeta, abstractmethod
 from xblock.plugin import default_select
@@ -447,6 +448,28 @@ class ModuleStoreReadBase(ModuleStoreRead):
         # default is to say yes by not raising an exception
         return {'default_impl': True}
 
+    @contextmanager
+    def default_store(self, store_type):
+        """
+        A context manager for temporarily changing the default store
+        """
+        if self.get_modulestore_type(None) != store_type:
+            raise ValueError(u"Cannot set default store to type {}".format(store_type))
+        yield
+
+    @contextmanager
+    def branch_setting(self, branch_setting, course_id=None):
+        """
+        A context manager for temporarily setting a store's branch value
+        """
+        previous_branch_setting_func = getattr(self, 'branch_setting_func', None)
+        try:
+            self.branch_setting_func = lambda: branch_setting
+            yield
+        finally:
+            self.branch_setting_func = previous_branch_setting_func
+
+
 class ModuleStoreWriteBase(ModuleStoreReadBase, ModuleStoreWrite):
     '''
     Implement interface functionality that can be shared.
@@ -530,6 +553,29 @@ class ModuleStoreWriteBase(ModuleStoreReadBase, ModuleStoreWrite):
         super(ModuleStoreWriteBase, self).clone_course(source_course_id, dest_course_id, user_id)
         return dest_course_id
 
+    @contextmanager
+    def bulk_write_operations(self, course_id):
+        """
+        A context manager for notifying the store of bulk write events.
+
+        In the case of Mongo, it temporarily disables refreshing the metadata inheritance tree
+        until the bulk operation is completed.
+        """
+        # TODO
+        # Make this multi-process-safe if future operations need it.
+        # Right now, only Import Course, Clone Course, and Delete Course use this, so
+        # it's ok if the cached metadata in the memcache is invalid when another
+        # request comes in for the same course.
+        try:
+            if hasattr(self, '_begin_bulk_write_operation'):
+                self._begin_bulk_write_operation(course_id)
+            yield
+        finally:
+            # check for the begin method here,
+            # since it's an error if an end method is not defined when a begin method is
+            if hasattr(self, '_begin_bulk_write_operation'):
+                self._end_bulk_write_operation(course_id)
+
 
 def only_xmodules(identifier, entry_points):
     """Only use entry_points that are supplied by the xmodule package"""

common/lib/xmodule/xmodule/modulestore/mixed.py

@@ -492,9 +492,9 @@ class MixedModuleStore(ModuleStoreWriteBase):
             raise NotImplementedError(u"Cannot call {} on store {}".format(method, store))
 
     @contextmanager
-    def set_default_store(self, store_type):
+    def default_store(self, store_type):
         """
-        A context manager for temporarily changing the default store in the Mixed modulestore
+        A context manager for temporarily changing the default store in the Mixed modulestore to the given store type
         """
         previous_store_list = self.modulestores
         found = False
@@ -509,50 +509,22 @@ class MixedModuleStore(ModuleStoreWriteBase):
         finally:
             self.modulestores = previous_store_list
 
-
-@contextmanager
-def store_branch_setting(store, branch_setting):
+    @contextmanager
+    def branch_setting(self, branch_setting, course_id=None):
         """
-    A context manager for temporarily setting a store's branch value
-
-    Note: to be effective, the store must be a direct pointer to the underlying store;
-        not the intermediary Mixed store.
+        A context manager for temporarily setting the branch value for the given course' store
+        to the given branch_setting.  If course_id is None, the default store is used.
         """
-    assert not isinstance(store, MixedModuleStore)
-
-    try:
-        previous_branch_setting_func = store.branch_setting_func
-        store.branch_setting_func = lambda: branch_setting
+        store = self._get_modulestore_for_courseid(course_id)
+        with store.branch_setting(branch_setting, course_id):
             yield
-    finally:
-        store.branch_setting_func = previous_branch_setting_func
 
-
-@contextmanager
-def store_bulk_write_operations_on_course(store, course_id):
+    @contextmanager
+    def bulk_write_operations(self, course_id):
         """
         A context manager for notifying the store of bulk write events.
-
-    In the case of Mongo, it temporarily disables refreshing the metadata inheritance tree
-    until the bulk operation is completed.
-
-    The store can be either the Mixed modulestore or a direct pointer to the underlying store.
+        If course_id is None, the default store is used.
         """
-
-    # TODO
-    # Make this multi-process-safe if future operations need it.
-    # Right now, only Import Course, Clone Course, and Delete Course use this, so
-    # it's ok if the cached metadata in the memcache is invalid when another
-    # request comes in for the same course.
-
-    # if the caller passed in the mixed modulestore, get a direct pointer to the underlying store
-    if hasattr(store, '_get_modulestore_by_course_id'):
-        store = store._get_modulestore_by_course_id(course_id)
-
-    try:
-        if hasattr(store, 'begin_bulk_write_operation_on_course'):
-            store.begin_bulk_write_operation_on_course(course_id)
+        store = self._get_modulestore_for_courseid(course_id)
+        with store.bulk_write_operations(course_id):
             yield
-    finally:
-        if hasattr(store, 'begin_bulk_write_operation_on_course'):
-            store.end_bulk_write_operation_on_course(course_id)

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

@@ -390,13 +390,13 @@ class MongoModuleStore(ModuleStoreWriteBase):
         self.ignore_write_events_on_courses = set()
         self._course_run_cache = {}
 
-    def begin_bulk_write_operation_on_course(self, course_id):
+    def _begin_bulk_write_operation(self, course_id):
         """
         Prevent updating the meta-data inheritance cache for the given course
         """
         self.ignore_write_events_on_courses.add(course_id)
 
-    def end_bulk_write_operation_on_course(self, course_id):
+    def _end_bulk_write_operation(self, course_id):
         """
         Restart updating the meta-data inheritance cache for the given course.
         Refresh the meta-data inheritance cache now since it was temporarily disabled.

common/lib/xmodule/xmodule/modulestore/xml_exporter.py

@@ -9,7 +9,6 @@ from xmodule.contentstore.content import StaticContent
 from xmodule.exceptions import NotFoundError
 from xmodule.modulestore import EdxJSONEncoder, ModuleStoreEnum
 from xmodule.modulestore.inheritance import own_metadata
-from xmodule.modulestore.mixed import store_branch_setting
 from fs.osfs import OSFS
 from json import dumps
 import json
@@ -45,7 +44,7 @@ def export_to_xml(modulestore, contentstore, course_key, root_dir, course_dir):
     root = lxml.etree.Element('unknown')
 
     # export only the published content
-    with store_branch_setting(course.runtime.modulestore, ModuleStoreEnum.Branch.published_only):
+    with modulestore.branch_setting(ModuleStoreEnum.Branch.published_only, course_key):
         course.add_xml_to_node(root)
 
     with export_fs.open('course.xml', 'w') as course_xml:

common/lib/xmodule/xmodule/modulestore/xml_importer.py

@@ -10,7 +10,6 @@ from xmodule.x_module import XModuleDescriptor
 from opaque_keys.edx.keys import UsageKey
 from xblock.fields import Scope, Reference, ReferenceList, ReferenceValueDict
 from xmodule.contentstore.content import StaticContent
-from xmodule.modulestore.mixed import store_bulk_write_operations_on_course
 from .inheritance import own_metadata
 from xmodule.errortracker import make_error_tracker
 from .store_utilities import rewrite_nonportable_content_links
@@ -179,7 +178,7 @@ def import_from_xml(
                 )
                 continue
 
-        with store_bulk_write_operations_on_course(store, dest_course_id):
+        with store.bulk_write_operations(dest_course_id):
             course_data_path = None
 
             if verbose: