Use BSON preconf stage and store response values under top-level keys, so created_at attribute is compatible with TTL index

4 files changed, 51 insertions, 37 deletions

diff --git a/requests_cache/backends/mongodb.py b/requests_cache/backends/mongodb.py
index 0fb83f2..2e4bf5a 100644
--- a/requests_cache/backends/mongodb.py
+++ b/requests_cache/backends/mongodb.py
@@ -78,19 +78,27 @@ API Reference
    :nosignatures:
 """
 from datetime import timedelta
-from typing import Iterable, Union
+from logging import getLogger
+from typing import Iterable, Mapping, Union
 
 from pymongo import MongoClient
 from pymongo.errors import OperationFailure
 
 from .._utils import get_valid_kwargs
 from ..expiration import NEVER_EXPIRE, get_expiration_seconds
-from ..serializers import dict_serializer
+from ..serializers import SerializerPipeline
+from ..serializers.preconf import bson_preconf_stage
 from . import BaseCache, BaseStorage
 
+document_serializer = SerializerPipeline([bson_preconf_stage], is_binary=False)
+logger = getLogger(__name__)
+
 
 # TODO: TTL tests
 # TODO: Example of viewing responses with MongoDB VSCode plugin or other GUI
+# TODO: Is there any reason to support custom serializers here?
+# TODO: Save items with different cache keys to avoid conflicts with old serialization format?
+# TODO: Set TTL for redirects? Or just clean up with remove_invalid_redirects()?
 class MongoCache(BaseCache):
     """MongoDB cache backend
 
@@ -153,21 +161,28 @@ class MongoDict(BaseStorage):
         if overwrite:
             try:
                 self.collection.drop_index('ttl_idx')
+                logger.info('Dropped TTL index')
             except OperationFailure:
                 pass
+
         ttl = get_expiration_seconds(ttl)
         if ttl and ttl != NEVER_EXPIRE:
+            logger.info(f'Creating TTL index for {ttl} seconds')
             self.collection.create_index('created_at', name='ttl_idx', expireAfterSeconds=ttl)
 
     def __getitem__(self, key):
         result = self.collection.find_one({'_id': key})
         if result is None:
             raise KeyError
-        return result['data']
+        return result['data'] if 'data' in result else result
 
     def __setitem__(self, key, item):
-        doc = {'_id': key, 'data': item}
-        self.collection.replace_one({'_id': key}, doc, upsert=True)
+        """If ``item`` is already a dict, its values will be stored under top-level keys.
+        Otherwise, it will be stored under a 'data' key.
+        """
+        if not isinstance(item, Mapping):
+            item = {'data': item}
+        self.collection.replace_one({'_id': key}, item, upsert=True)
 
     def __delitem__(self, key):
         result = self.collection.find_one_and_delete({'_id': key}, {'_id': True})
@@ -192,30 +207,14 @@ class MongoDict(BaseStorage):
 class MongoPickleDict(MongoDict):
     """Same as :class:`MongoDict`, but serializes values before saving.
 
-    By default, responses are only partially serialized (unstructured into a dict), and stored as a
-    document.
+    By default, responses are only partially serialized into a MongoDB-compatible document mapping.
     """
 
-    def __init__(
-        self,
-        db_name: str,
-        collection_name: str = 'http_cache',
-        connection: MongoClient = None,
-        ttl: int = None,
-        serializer=None,
-        **kwargs,
-    ):
-        super().__init__(
-            db_name,
-            collection_name=collection_name,
-            connection=connection,
-            ttl=ttl,
-            serializer=serializer or dict_serializer,
-            **kwargs,
-        )
-
-    def __setitem__(self, key, item):
-        super().__setitem__(key, self.serializer.dumps(item))
+    def __init__(self, *args, serializer=None, **kwargs):
+        super().__init__(*args, serializer=serializer or document_serializer, **kwargs)
 
     def __getitem__(self, key):
         return self.serializer.loads(super().__getitem__(key))
+
+    def __setitem__(self, key, item):
+        super().__setitem__(key, self.serializer.dumps(item))
diff --git a/requests_cache/serializers/__init__.py b/requests_cache/serializers/__init__.py
index dc78489..91c5e4c 100644
--- a/requests_cache/serializers/__init__.py
+++ b/requests_cache/serializers/__init__.py
@@ -19,6 +19,7 @@ __all__ = [
     'SerializerPipeline',
     'Stage',
     'bson_serializer',
+    'dict_serializer',
     'json_serializer',
     'pickle_serializer',
     'safe_pickle_serializer',
@@ -29,7 +30,6 @@ __all__ = [
 
 SERIALIZERS = {
     'bson': bson_serializer,
-    'dict': dict_serializer,
     'json': json_serializer,
     'pickle': pickle_serializer,
     'yaml': yaml_serializer,
diff --git a/requests_cache/serializers/cattrs.py b/requests_cache/serializers/cattrs.py
index e0bde94..da7fd79 100644
--- a/requests_cache/serializers/cattrs.py
+++ b/requests_cache/serializers/cattrs.py
@@ -28,8 +28,8 @@ class CattrStage(Stage):
     on its own, or as a stage within a :py:class:`.SerializerPipeline`.
     """
 
-    def __init__(self, factory: Callable[..., GenConverter] = None):
-        self.converter = init_converter(factory)
+    def __init__(self, factory: Callable[..., GenConverter] = None, **kwargs):
+        self.converter = init_converter(factory, **kwargs)
 
     def dumps(self, value: CachedResponse) -> Dict:
         if not isinstance(value, CachedResponse):
@@ -42,14 +42,22 @@ class CattrStage(Stage):
         return self.converter.structure(value, cl=CachedResponse)
 
 
-def init_converter(factory: Callable[..., GenConverter] = None):
-    """Make a converter to structure and unstructure nested objects within a :py:class:`.CachedResponse`"""
+def init_converter(factory: Callable[..., GenConverter] = None, convert_datetime: bool = True):
+    """Make a converter to structure and unstructure nested objects within a
+    :py:class:`.CachedResponse`
+
+    Args:
+        factory: An optional factory function that returns a ``cattrs`` converter
+        convert_datetime: May be set to ``False`` for pre-configured converters that already have
+            datetime support
+    """
     factory = factory or GenConverter
     converter = factory(omit_if_default=True)
 
     # Convert datetimes to and from iso-formatted strings
-    converter.register_unstructure_hook(datetime, lambda obj: obj.isoformat() if obj else None)  # type: ignore
-    converter.register_structure_hook(datetime, _to_datetime)
+    if convert_datetime:
+        converter.register_unstructure_hook(datetime, lambda obj: obj.isoformat() if obj else None)  # type: ignore
+        converter.register_structure_hook(datetime, _to_datetime)
 
     # Convert timedeltas to and from float values in seconds
     converter.register_unstructure_hook(timedelta, lambda obj: obj.total_seconds() if obj else None)  # type: ignore
@@ -66,6 +74,10 @@ def init_converter(factory: Callable[..., GenConverter] = None):
     converter.register_structure_hook(HTTPHeaderDict, lambda obj, cls: HTTPHeaderDict(obj))
 
     # Tell cattrs to resolve forward references (required for CachedResponse.history)
+    converter.register_unstructure_hook_func(
+        lambda cls: cls.__class__ is ForwardRef,
+        lambda obj, cls=None: converter.unstructure(obj, cls.__forward_value__ if cls else None),
+    )
     converter.register_structure_hook_func(
         lambda cls: cls.__class__ is ForwardRef,
         lambda obj, cls: converter.structure(obj, cls.__forward_value__),
diff --git a/requests_cache/serializers/preconf.py b/requests_cache/serializers/preconf.py
index 1236c44..b0ad069 100644
--- a/requests_cache/serializers/preconf.py
+++ b/requests_cache/serializers/preconf.py
@@ -22,10 +22,11 @@ from .cattrs import CattrStage
 from .pipeline import SerializerPipeline, Stage
 
 
-def make_stage(preconf_module: str):
+def make_stage(preconf_module: str, **kwargs):
     """Create a preconf serializer stage from a module name, if dependencies are installed"""
     try:
-        return CattrStage(import_module(preconf_module).make_converter)
+        factory = import_module(preconf_module).make_converter
+        return CattrStage(factory, **kwargs)
     except ImportError as e:
         return get_placeholder_class(e)
 
@@ -33,7 +34,9 @@ def make_stage(preconf_module: str):
 # Pre-serialization stages
 base_stage = CattrStage()  #: Base stage for all serializer pipelines
 utf8_encoder = Stage(dumps=str.encode, loads=lambda x: x.decode())  #: Encode to bytes
-bson_preconf_stage = make_stage('cattr.preconf.bson')  #: Pre-serialization steps for BSON
+bson_preconf_stage = make_stage(
+    'cattr.preconf.bson', convert_datetime=False
+)  #: Pre-serialization steps for BSON
 json_preconf_stage = make_stage('cattr.preconf.json')  #: Pre-serialization steps for JSON
 msgpack_preconf_stage = make_stage('cattr.preconf.msgpack')  #: Pre-serialization steps for msgpack
 orjson_preconf_stage = make_stage('cattr.preconf.orjson')  #: Pre-serialization steps for orjson