Handle CorruptGridFile error in GridFS backend

requests_cache/backends/gridfs.py

@@ -15,7 +15,7 @@
 from threading import RLock
 
 from gridfs import GridFS
-from gridfs.errors import FileExists
+from gridfs.errors import CorruptGridFile, FileExists
 from pymongo import MongoClient
 
 from .._utils import get_valid_kwargs
@@ -69,11 +69,15 @@ def __init__(self, db_name, collection_name=None, connection=None, **kwargs):
         self._lock = RLock()
 
     def __getitem__(self, key):
-        with self._lock:
-            result = self.fs.find_one({'_id': key})
-            if result is None:
-                raise KeyError
-            return self.serializer.loads(result.read())
+        try:
+            with self._lock:
+                result = self.fs.find_one({'_id': key})
+                if result is None:
+                    raise KeyError
+                return self.serializer.loads(result.read())
+        except CorruptGridFile as e:
+            logger.warning(e, exc_info=True)
+            raise KeyError
 
     def __setitem__(self, key, item):
         value = self.serializer.dumps(item)
@@ -83,6 +87,7 @@ def __setitem__(self, key, item):
             try:
                 self.fs.delete(key)
                 self.fs.put(value, encoding=encoding, **{'_id': key})
+            # This can happen because GridFS is not thread-safe for concurrent writes
             except FileExists as e:
                 logger.warning(e, exc_info=True)
 

tests/integration/base_cache_test.py

@@ -353,6 +353,7 @@ def test_concurrency(self, iteration, executor_class):
         """
         start = time()
         url = httpbin('anything')
+        self.init_session(clear=True)
 
         session_factory = partial(self.init_session, clear=False)
         request_func = partial(_send_request, session_factory, url)

tests/integration/test_mongodb.py

@@ -1,6 +1,8 @@
 from unittest.mock import patch
 
 import pytest
+from gridfs import GridFS
+from gridfs.errors import CorruptGridFile, FileExists
 from pymongo import MongoClient
 
 from requests_cache._utils import get_valid_kwargs
@@ -65,6 +67,22 @@ def test_connection_kwargs(self, mock_get_valid_kwargs, mock_client, mock_gridfs
         GridFSPickleDict('test', host='http://0.0.0.0', port=1234, invalid_kwarg='???')
         mock_client.assert_called_with(host='http://0.0.0.0', port=1234)
 
+    def test_corrupt_file(self):
+        """A corrupted file should be handled and raise a KeyError instead"""
+        cache = self.init_cache()
+        cache['key'] = 'value'
+        with pytest.raises(KeyError), patch.object(GridFS, 'find_one', side_effect=CorruptGridFile):
+            cache['key']
+
+    def test_file_exists(self):
+        cache = self.init_cache()
+
+        # This write should just quiety fail
+        with patch.object(GridFS, 'put', side_effect=FileExists):
+            cache['key'] = 'value_1'
+
+        assert 'key' not in cache
+
 
 class TestGridFSCache(BaseCacheTest):
     backend_class = GridFSCache