refactor: replace MongoDBProxy usage with direct pymongo config

Inspired by openedx#35213 Closes openedx#35210 Co-Authored-By: Yash Athwani <[email protected]>
kdmccormick · Sep 27, 2024 · 5e139ce · 5e139ce
1 parent de7fa09
commit 5e139ce
Show file tree

Hide file tree

Showing 11 changed files with 17 additions and 53 deletions.
diff --git a/requirements/edx/base.txt b/requirements/edx/base.txt
@@ -828,8 +828,6 @@ openedx-learning==0.13.1
     # via
     #   -c requirements/edx/../constraints.txt
     #   -r requirements/edx/kernel.in
-openedx-mongodbproxy==0.2.1
-    # via -r requirements/edx/kernel.in
 optimizely-sdk==4.1.1
     # via
     #   -c requirements/edx/../constraints.txt
@@ -955,7 +953,6 @@ pymongo==4.4.0
     #   edx-opaque-keys
     #   event-tracking
     #   mongoengine
-    #   openedx-mongodbproxy
 pynacl==1.5.0
     # via
     #   edx-django-utils

diff --git a/requirements/edx/development.txt b/requirements/edx/development.txt
@@ -1378,10 +1378,6 @@ openedx-learning==0.13.1
     #   -c requirements/edx/../constraints.txt
     #   -r requirements/edx/doc.txt
     #   -r requirements/edx/testing.txt
-openedx-mongodbproxy==0.2.1
-    # via
-    #   -r requirements/edx/doc.txt
-    #   -r requirements/edx/testing.txt
 optimizely-sdk==4.1.1
     # via
     #   -c requirements/edx/../constraints.txt
@@ -1637,7 +1633,6 @@ pymongo==4.4.0
     #   edx-opaque-keys
     #   event-tracking
     #   mongoengine
-    #   openedx-mongodbproxy
 pynacl==1.5.0
     # via
     #   -r requirements/edx/doc.txt

diff --git a/requirements/edx/doc.txt b/requirements/edx/doc.txt
@@ -987,8 +987,6 @@ openedx-learning==0.13.1
     # via
     #   -c requirements/edx/../constraints.txt
     #   -r requirements/edx/base.txt
-openedx-mongodbproxy==0.2.1
-    # via -r requirements/edx/base.txt
 optimizely-sdk==4.1.1
     # via
     #   -c requirements/edx/../constraints.txt
@@ -1149,7 +1147,6 @@ pymongo==4.4.0
     #   edx-opaque-keys
     #   event-tracking
     #   mongoengine
-    #   openedx-mongodbproxy
 pynacl==1.5.0
     # via
     #   -r requirements/edx/base.txt

diff --git a/requirements/edx/kernel.in b/requirements/edx/kernel.in
@@ -120,7 +120,6 @@ openedx-django-require
 openedx-events                      # Open edX Events from Hooks Extension Framework (OEP-50)
 openedx-filters                     # Open edX Filters from Hooks Extension Framework (OEP-50)
 openedx-learning                    # Open edX Learning core (experimental)
-openedx-mongodbproxy
 openedx-django-wiki
 path
 piexif                              # Exif image metadata manipulation, used in the profile_images app

diff --git a/requirements/edx/testing.txt b/requirements/edx/testing.txt
@@ -1038,8 +1038,6 @@ openedx-learning==0.13.1
     # via
     #   -c requirements/edx/../constraints.txt
     #   -r requirements/edx/base.txt
-openedx-mongodbproxy==0.2.1
-    # via -r requirements/edx/base.txt
 optimizely-sdk==4.1.1
     # via
     #   -c requirements/edx/../constraints.txt
@@ -1234,7 +1232,6 @@ pymongo==4.4.0
     #   edx-opaque-keys
     #   event-tracking
     #   mongoengine
-    #   openedx-mongodbproxy
 pynacl==1.5.0
     # via
     #   -r requirements/edx/base.txt

diff --git a/xmodule/contentstore/mongo.py b/xmodule/contentstore/mongo.py
@@ -12,7 +12,7 @@
 from bson.son import SON
 from fs.osfs import OSFS
 from gridfs.errors import NoFile, FileExists
-from mongodb_proxy import autoretry_read
+
 from opaque_keys.edx.keys import AssetKey
 
 from xmodule.contentstore.content import XASSET_LOCATION_TAG
@@ -39,8 +39,7 @@ def __init__(
         :param collection: ignores but provided for consistency w/ other doc_store_config patterns
         """
         # GridFS will throw an exception if the Database is wrapped in a MongoProxy. So don't wrap it.
-        # The appropriate methods below are marked as autoretry_read - those methods will handle
-        # the AutoReconnect errors.
+
         self.connection_params = {
             'db': db,
             'host': host,
@@ -164,7 +163,6 @@ def delete(self, location_or_id):
         # Deletes of non-existent files are considered successful
         self.fs.delete(location_or_id)
 
-    @autoretry_read()
     def find(self, location, throw_on_not_found=True, as_stream=False):  # lint-amnesty, pylint: disable=arguments-differ
         content_id, __ = self.asset_db_key(location)
 
@@ -292,7 +290,7 @@ def remove_redundant_content_for_courses(self):
             self.fs_files.remove(query)
         return assets_to_delete
 
-    @autoretry_read()
+
     def _get_all_content_for_course(self,
                                     course_key,
                                     get_thumbnails=False,
@@ -424,7 +422,6 @@ def set_attrs(self, location, attr_dict):
         if result.matched_count == 0:
             raise NotFoundError(asset_db_key)
 
-    @autoretry_read()
     def get_attrs(self, location):
         """
         Gets all of the attributes associated with the given asset. Note, returns even built in attrs

diff --git a/xmodule/modulestore/mongo/base.py b/xmodule/modulestore/mongo/base.py
@@ -23,7 +23,6 @@
 import pymongo
 from bson.son import SON
 from fs.osfs import OSFS
-from mongodb_proxy import autoretry_read
 from opaque_keys.edx.keys import CourseKey, UsageKey
 from opaque_keys.edx.locator import BlockUsageLocator, CourseLocator, LibraryLocator
 from path import Path as path
@@ -578,7 +577,6 @@ def _drop_database(self, database=True, collections=True, connections=True):
         if connections:
             connection.close()
 
-    @autoretry_read()
     def fill_in_run(self, course_key):
         """
         In mongo some course_keys are used without runs. This helper function returns
@@ -737,7 +735,6 @@ def _load_items(self, course_key, items, using_descriptor_system=None, for_paren
             for item in items
         ]
 
-    @autoretry_read()
     def get_course_summaries(self, **kwargs):
         """
         Returns a list of `CourseSummary`. This accepts an optional parameter of 'org' which
@@ -786,7 +783,6 @@ def extract_course_summary(course):
 
         return courses_summaries
 
-    @autoretry_read()
     def get_courses(self, **kwargs):
         '''
         Returns a list of course descriptors. This accepts an optional parameter of 'org' which
@@ -821,7 +817,6 @@ def get_courses(self, **kwargs):
         )
         return [course for course in base_list if not isinstance(course, ErrorBlock)]
 
-    @autoretry_read()
     def _find_one(self, location):
         '''Look for a given location in the collection. If the item is not present, raise
         ItemNotFoundError.
@@ -867,7 +862,6 @@ def get_course(self, course_key, **kwargs):  # lint-amnesty, pylint: disable=arg
         except ItemNotFoundError:
             return None
 
-    @autoretry_read()
     def has_course(self, course_key, ignore_case=False, **kwargs):  # lint-amnesty, pylint: disable=arguments-differ
         """
         Returns the course_id of the course if it was found, else None
@@ -962,7 +956,6 @@ def _id_dict_to_son(id_dict):
             for key in ('tag', 'org', 'course', 'category', 'name', 'revision')
         ])
 
-    @autoretry_read()
     def get_items(  # lint-amnesty, pylint: disable=arguments-differ
             self,
             course_id,

diff --git a/xmodule/modulestore/split_mongo/mongo_connection.py b/xmodule/modulestore/split_mongo/mongo_connection.py
@@ -17,7 +17,6 @@
 from django.db.transaction import TransactionManagementError
 import pymongo
 import pytz
-from mongodb_proxy import autoretry_read
 # Import this just to export it
 from pymongo.errors import DuplicateKeyError  # pylint: disable=unused-import
 from edx_django_utils import monitoring
@@ -363,7 +362,6 @@ def get_structure(self, key, course_context=None):
 
             return structure
 
-    @autoretry_read()
     def find_structures_by_id(self, ids, course_context=None):
         """
         Return all structures that specified in ``ids``.
@@ -380,7 +378,6 @@ def find_structures_by_id(self, ids, course_context=None):
             tagger.measure("structures", len(docs))
             return docs
 
-    @autoretry_read()
     def find_courselike_blocks_by_id(self, ids, block_type, course_context=None):
         """
         Find all structures that specified in `ids`. Among the blocks only return block whose type is `block_type`.

diff --git a/xmodule/modulestore/split_mongo/split.py b/xmodule/modulestore/split_mongo/split.py
@@ -63,7 +63,6 @@
 
 from bson.objectid import ObjectId
 from ccx_keys.locator import CCXBlockUsageLocator, CCXLocator
-from mongodb_proxy import autoretry_read
 from opaque_keys.edx.keys import CourseKey
 from opaque_keys.edx.locator import (
     BlockUsageLocator,
@@ -953,7 +952,6 @@ def _create_library_locator(self, library_info, branch):
             branch=branch,
         )
 
-    @autoretry_read()
     def get_courses(self, branch, **kwargs):  # lint-amnesty, pylint: disable=arguments-differ
         """
         Returns a list of course blocks matching any given qualifiers.
@@ -969,7 +967,6 @@ def get_courses(self, branch, **kwargs):  # lint-amnesty, pylint: disable=argume
         # get the blocks for each course index (s/b the root)
         return self._get_structures_for_branch_and_locator(branch, self._create_course_locator, **kwargs)
 
-    @autoretry_read()
     def get_course_summaries(self, branch, **kwargs):
         """
         Returns a list of `CourseSummary` which matching any given qualifiers.
@@ -1026,7 +1023,6 @@ def get_library_keys(self):
             in self.find_matching_course_indexes(branch="library")
         })
 
-    @autoretry_read()
     def get_library_summaries(self, **kwargs):
         """
         Returns a list of `LibrarySummary` objects.
@@ -3255,7 +3251,6 @@ def _update_block_in_structure(self, structure, block_key, content):
         """
         structure['blocks'][block_key] = content
 
-    @autoretry_read()
     def find_courses_by_search_target(self, field_name, field_value):
         """
         Find all the courses which cached that they have the given field with the given value.

diff --git a/xmodule/modulestore/tests/test_split_mongo_mongo_connection.py b/xmodule/modulestore/tests/test_split_mongo_mongo_connection.py
@@ -16,11 +16,17 @@ class TestHeartbeatFailureException(unittest.TestCase):
 
     @patch('pymongo.MongoClient')
     @patch('pymongo.database.Database')
-    def test_heartbeat_raises_exception_when_connection_alive_is_false(self, *calls):
-        # pylint: disable=W0613
-        with patch('mongodb_proxy.MongoProxy') as mock_proxy:
-            mock_proxy.return_value.admin.command.side_effect = ConnectionFailure('Test')
-            useless_conn = MongoPersistenceBackend('useless', 'useless', 'useless')
-
-            with pytest.raises(HeartbeatFailure):
-                useless_conn.heartbeat()
+    def test_heartbeat_retries_on_failure(self, MockDatabase, MockClient):
+        # Setup mock client and database
+        mock_client = MockClient.return_value
+        mock_database = MockDatabase.return_value
+        mock_database.admin.command.side_effect = ConnectionFailure('Test')
+
+        useless_conn = MongoPersistenceBackend('useless', 'useless', 'useless')
+
+        # Verify that the heartbeat method raises a HeartbeatFailure
+        with pytest.raises(HeartbeatFailure):
+            useless_conn.heartbeat()
+
+        # Assert that retries are handled correctly
+        self.assertGreater(mock_database.admin.command.call_count, 1)  # Ensure retries happened
diff --git a/xmodule/mongo_utils.py b/xmodule/mongo_utils.py
@@ -6,7 +6,6 @@
 import logging
 
 import pymongo
-from mongodb_proxy import MongoProxy
 from pymongo.read_preferences import (  # lint-amnesty, pylint: disable=unused-import
     ReadPreference,
     _MONGOS_MODES,
@@ -69,14 +68,6 @@ def connect_to_mongodb(
         connection_params.update({'username': user, 'password': password, 'authSource': auth_source})
 
     mongo_conn = pymongo.MongoClient(**connection_params)
-
-    if proxy:
-        mongo_conn = MongoProxy(
-            mongo_conn[db],
-            wait_time=retry_wait_time
-        )
-        return mongo_conn
-
     return mongo_conn[db]