CASB-178258: Changes for pymongo-4.7+

django_mongodb_engine/base.py

@@ -73,7 +73,7 @@ def sql_flush(self, style, tables, sequence_list, allow_cascade=False):
             options = collection.options()
 
             if not options.get('capped', False):
-                collection.remove({})
+                collection.delete_many({})
 
         return []
 
@@ -175,7 +175,7 @@ class DatabaseValidation(NonrelDatabaseValidation):
 class DatabaseIntrospection(NonrelDatabaseIntrospection):
 
     def table_names(self, cursor=None):
-        return self.connection.database.collection_names()
+        return self.connection.database.list_collection_names()
 
     def sequence_list(self):
         # Only required for backends that use integer primary keys.
@@ -203,7 +203,7 @@ def __init__(self, *args, **kwargs):
 
     def get_collection(self, name, **kwargs):
         if (kwargs.pop('existing', False) and
-                name not in self.database.collection_names()):
+                name not in self.database.list_collection_names()):
             return None
         collection = self.collection_class(self.database, name, **kwargs)
         if settings.DEBUG:
@@ -272,7 +272,7 @@ def pop(name, default=None):
 
         # In PyMongo3.6.0, MongoClient is asynchronous. To have consistent behaviour, making a cheap query
         try:
-            self.database.command('ismaster')
+            self.database.command('hello')
         except Exception:
             exc_info = sys.exc_info()
             raise ImproperlyConfigured(exc_info[1]).with_traceback(exc_info[2])

django_mongodb_engine/compiler.py

@@ -109,10 +109,13 @@ def fetch(self, low_mark, high_mark):
 
     @safe_call
     def count(self, limit=None):
-        results = self.get_cursor()
-        if limit is not None:
-            results.limit(limit)
-        return results.count()
+        if(self.mongo_query):
+             return self.collection.count_documents(self.mongo_query)
+        return self.collection.estimated_document_count()
+        # results = self.get_cursor()
+        # if limit is not None:
+        #     results.limit(limit)
+        # return results.count()
 
     @safe_call
     def order_by(self, ordering):
@@ -129,7 +132,7 @@ def order_by(self, ordering):
     @safe_call
     def delete(self):
         options = self.connection.operation_flags.get('delete', {})
-        self.collection.remove(self.mongo_query, **options)
+        self.collection.delete_many(self.mongo_query, **options)
 
     def get_cursor(self):
         if self.query.low_mark == self.query.high_mark:
@@ -386,11 +389,15 @@ def insert(self, docs, return_id=False):
 
         collection = self.get_collection()
         options = self.connection.operation_flags.get('save', {})
+        if '_id' in doc:
+            collection.replace_one({'_id': doc['_id']}, doc, upsert=True, **options)
+            return_doc_id = doc['_id']
+        else:
+            result = collection.insert_one(doc, **options)
+            return_doc_id = result.inserted_id
 
         if return_id:
-            return collection.save(doc, **options)
-        else:
-            collection.save(doc, **options)
+            return return_doc_id
 
 
 # TODO: Define a common nonrel API for updates and add it to the nonrel
@@ -437,9 +444,12 @@ def execute_update(self, update_spec, multi=True, **kwargs):
             return 0
         options = self.connection.operation_flags.get('update', {})
         options = dict(options, **kwargs)
-        info = collection.update(criteria, update_spec, multi=multi, **options)
+        if multi:
+            info = collection.update_many(criteria, update_spec, **options)
+        else:
+            info = collection.update_one(criteria, update_spec, **options)
         if info is not None:
-            return info.get('n')
+            return info.raw_result['n']
 
 
 class SQLDeleteCompiler(NonrelDeleteCompiler, SQLCompiler):

django_mongodb_engine/creation.py

@@ -44,13 +44,14 @@ def ensure_index(*args, **kwargs):
                       (meta.app_label, meta.object_name))
                 ensure_index.first_index = False
             try:
-                return collection.ensure_index(*args, **kwargs)
+                return collection.create_index(*args, **kwargs)
             except OperationFailure as e:
-                # Try with a short name for the index in case the
-                # auto-generated index name is too long
-                index_name = meta.object_name + "_index_" + str(randint(1,100000))
-                print("Error installing index on %s: %s, trying shorter index name: %s" % (meta.object_name, str(e), index_name))
-                return collection.ensure_index(*args, name=index_name, **kwargs)
+                if 'key too large to index' in str(e):
+                    # Try with a short name for the index in case the
+                    # auto-generated index name is too long
+                    index_name = meta.object_name + "_index_" + str(randint(1,100000))
+                    print("Error installing index on %s: %s, trying shorter index name: %s" % (meta.object_name, str(e), index_name))
+                    return collection.create_index(*args, name=index_name, **kwargs)
 
         ensure_index.first_index = True
 

django_mongodb_engine/south_adapter.py

@@ -42,7 +42,7 @@ def add_column(self, table_name, field_name, field, keep_default=True):
                 db_prep_save = field.get_db_prep_save(default, connection=connection)
                 default = connection.ops.value_for_db(db_prep_save, field)
                 # Update all the documents that haven't got this field yet
-                collection.update({name: {'$exists': False}},
+                collection.update_one({name: {'$exists': False}},
                                   {'$set': {name: default}})
             if not keep_default:
                 field.default = NOT_PROVIDED
@@ -53,11 +53,11 @@ def alter_column(self, table_name, column_name, field, explicit_name=True):
 
     def delete_column(self, table_name, name):
         collection = self._get_collection(table_name)
-        collection.update({}, {'$unset': {name: 1}})
+        collection.update_one({}, {'$unset': {name: 1}})
 
     def rename_column(self, table_name, old, new):
         collection = self._get_collection(table_name)
-        collection.update({}, {'$rename': {old: new}})
+        collection.update_one({}, {'$rename': {old: new}})
 
     def create_unique(self, table_name, columns, drop_dups=False):
         collection = self._get_collection(table_name)

django_mongodb_engine/storage.py

@@ -13,7 +13,7 @@ def _get_subcollections(collection):
     Returns all sub-collections of `collection`.
     """
     # XXX: Use the MongoDB API for this once it exists.
-    for name in collection.database.collection_names():
+    for name in collection.database.list_collection_names():
         cleaned = name[:name.rfind('.')]
         if cleaned != collection.name and cleaned.startswith(collection.name):
             yield cleaned

tests/storage/tests.py

@@ -19,7 +19,7 @@ def setUp(self):
 
     def tearDown(self):
         if hasattr(self.storage, '_db'):
-            for collection in self.storage._db.collection_names():
+            for collection in self.storage._db.list_collection_names():
                 if not collection.startswith('system.'):
                     self.storage._db.drop_collection(collection)