storage: add automatic expiry to MongoDB

This changes the record layout in MongoDB so that the modification time
is now stored as a datetime object instead of a timestamp. It enables us
to create an auto-expiry index on that field.

Note that expiring objects is an asynchronous task in MongoDB. Unlike
Redis, MongoDB uses collection-level expiry settings, therefore it is
not possible to create multiple wrapper objects that operate on the same
collection with different TTL values.

Additionally, there is a minor change that allows specifying the
database name in the URI instead of an argument. In case the name is
specified both in the URI and the argument, the one in the argument is
silently ignored.

Author: bajnokk

src/pyop/storage.py

@@ -5,7 +5,7 @@
 import json
 import pymongo
 from redis.client import Redis
-from time import time
+from datetime import datetime
 
 
 class StorageBase(ABC):
@@ -54,18 +54,29 @@ def ttl(self):
 
 
 class MongoWrapper(StorageBase):
-    def __init__(self, db_uri, db_name, collection):
+    def __init__(self, db_uri, db_name, collection, ttl=None):
         self._db_uri = db_uri
         self._coll_name = collection
         self._db = MongoDB(db_uri, db_name=db_name)
         self._coll = self._db.get_collection(collection)
         self._coll.create_index('lookup_key', unique=True)
 
+        if ttl is None or (isinstance(ttl, int) and ttl >= 0):
+            self._ttl = ttl
+        else:
+            raise ValueError("TTL must be a non-negative integer or None")
+        if ttl is not None:
+            self._coll.create_index(
+                'last_modified',
+                expireAfterSeconds=ttl,
+                name="expiry"
+            )
+
     def __setitem__(self, key, value):
         doc = {
             'lookup_key': key,
             'data': value,
-            'modified_ts': time()
+            'last_modified': datetime.utcnow()
         }
         self._coll.replace_one({'lookup_key': key}, doc, upsert=True)
 
@@ -143,14 +154,16 @@ def __init__(self, db_uri, db_name=None,
         if db_uri is None:
             raise ValueError('db_uri not supplied')
 
-        self._db_uri = db_uri
-        self._database_name = db_name
         self._sanitized_uri = None
 
         self._parsed_uri = pymongo.uri_parser.parse_uri(db_uri)
 
-        if self._parsed_uri.get('database') is None:
-            self._parsed_uri['database'] = db_name
+        db_name = self._parsed_uri.get('database') or db_name
+        if db_name is None:
+            raise ValueError(
+                "Database name must be provided either in the URI or as an argument"
+            )
+        self._database_name = self._parsed_uri['database'] = db_name
 
         if 'replicaSet' in kwargs and kwargs['replicaSet'] is None:
             del kwargs['replicaSet']

tests/pyop/test_storage.py

@@ -168,3 +168,15 @@ def new_time():
 
     def test_ttl(self):
         self.execute_ttl_test("redis://localhost/0", 3600)
+
+
+class TestMongoTTL(StorageTTLTest):
+    def set_time(self, offset, monkeypatch):
+        now = datetime.datetime.utcnow()
+        def new_time():
+            return now + datetime.timedelta(seconds=offset)
+
+        monkeypatch.setattr(mongomock, "utcnow", new_time)
+
+    def test_ttl(self):
+        self.execute_ttl_test("mongodb://localhost/pyop", 3600)