WIP: Minor workaround in MongoDB backend to support mongitamongita

2 files changed, 32 insertions, 4 deletions

diff --git a/mongita_test.py b/mongita_test.py
new file mode 100644
index 0000000..18a4864
--- /dev/null
+++ b/mongita_test.py
@@ -0,0 +1,24 @@
+from mongita import MongitaClientDisk
+from pymongo import MongoClient
+
+
+def test_ids(client):
+    collection = client['test_db']['test_collection']
+    collection.replace_one(
+        {'_id': 'id_from_filter'},
+        replacement={'key': 'value'},
+        upsert=True,
+    )
+    doc = collection.find_one({'_id': 'id_from_filter'})
+    print(f'Fetched document by ID: {doc}')
+
+    print('All IDs:')
+    for d in collection.find({}):
+        print(d['_id'])
+
+
+print('pymongo\n----------')
+test_ids(MongoClient())
+
+print('\nmongita\n----------')
+test_ids(MongitaClientDisk())
diff --git a/requests_cache/backends/mongodb.py b/requests_cache/backends/mongodb.py
index 76b33be..c240eb1 100644
--- a/requests_cache/backends/mongodb.py
+++ b/requests_cache/backends/mongodb.py
@@ -129,18 +129,22 @@ class MongoDict(BaseStorage):
         value = self.serialize(value)
         if not isinstance(value, Mapping):
             value = {'data': value}
-        self.collection.replace_one({'_id': key}, value, upsert=True)
+        self.collection.replace_one({'_id': key}, replacement=value, upsert=True)
 
     def __delitem__(self, key):
-        result = self.collection.find_one_and_delete({'_id': key}, {'_id': True})
+        result = self.collection.find_one_and_delete({'_id': key}, projection={'_id': True})
         if result is None:
             raise KeyError
 
     def __len__(self):
-        return self.collection.estimated_document_count()
+        # First attempt a fast estimated count, then fallback to full count (compat with mongita)
+        try:
+            return self.collection.estimated_document_count()
+        except NotImplementedError:
+            return self.collection.count_documents({})
 
     def __iter__(self):
-        for d in self.collection.find({}, {'_id': True}):
+        for d in self.collection.find({}, projection={'_id': True}):
             yield d['_id']
 
     def bulk_delete(self, keys: Iterable[str]):