1 import datetime 2 import logging 3 import typing as _t 4 5 from cachelib.base import BaseCache 6 from cachelib.serializers import BaseSerializer 7 8 9 class MongoDbCache(BaseCache): 10 """ 11 Implementation of cachelib.BaseCache that uses mongodb collection 12 as the backend. 13 14 Limitations: maximum MongoDB document size is 16mb 15 16 :param client: mongodb client or connection string 17 :param db: mongodb database name 18 :param collection: mongodb collection name 19 :param default_timeout: Set the timeout in seconds after which cache entries 20 expire 21 :param key_prefix: A prefix that should be added to all keys. 22 23 """ 24 25 serializer = BaseSerializer() 26 27 def __init__( 28 self, 29 client: _t.Any = None, 30 db: _t.Optional[str] = "cache-db", 31 collection: _t.Optional[str] = "cache-collection", 32 default_timeout: int = 300, 33 key_prefix: _t.Optional[str] = None, 34 **kwargs: _t.Any 35 ): 36 super().__init__(default_timeout) 37 try: 38 import pymongo # type: ignore 39 except ImportError: 40 logging.warning("no pymongo module found") 41 42 if client is None or isinstance(client, str): 43 client = pymongo.MongoClient(host=client) 44 self.client = client[db][collection] 45 index_info = self.client.index_information() 46 all_keys = { 47 subkey[0] for value in index_info.values() for subkey in value["key"] 48 } 49 if "id" not in all_keys: 50 self.client.create_index("id", unique=True) 51 self.key_prefix = key_prefix or "" 52 self.collection = collection 53 54 def _utcnow(self) -> _t.Any: 55 """Return a tz-aware UTC datetime representing the current time""" 56 return datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc) 57 58 def _expire_records(self) -> _t.Any: 59 res = self.client.delete_many({"expiration": {"$lte": self._utcnow()}}) 60 return res 61 62 def get(self, key: str) -> _t.Any: 63 """ 64 Get a cache item 65 66 :param key: The cache key of the item to fetch 67 :return: cache value if not expired, else None 68 """ 69 self._expire_records() 70 record = self.client.find_one({"id": self.key_prefix + key}) 71 value = None 72 if record: 73 value = self.serializer.loads(record["val"]) 74 return value 75 76 def delete(self, key: str) -> bool: 77 """ 78 Deletes an item from the cache. This is a no-op if the item doesn't 79 exist 80 81 :param key: Key of the item to delete. 82 :return: True if the key existed and was deleted 83 """ 84 res = self.client.delete_one({"id": self.key_prefix + key}) 85 deleted = bool(res.deleted_count > 0) 86 return deleted 87 88 def _set( 89 self, 90 key: str, 91 value: _t.Any, 92 timeout: _t.Optional[int] = None, 93 overwrite: _t.Optional[bool] = True, 94 ) -> _t.Any: 95 """ 96 Store a cache item, with the option to not overwrite existing items 97 98 :param key: Cache key to use 99 :param value: a serializable object 100 :param timeout: The timeout in seconds for the cached item, to override 101 the default 102 :param overwrite: If true, overwrite any existing cache item with key. 103 If false, the new value will only be stored if no 104 non-expired cache item exists with key. 105 :return: True if the new item was stored. 106 """ 107 timeout = self._normalize_timeout(timeout) 108 now = self._utcnow() 109 110 if not overwrite: 111 # fail if a non-expired item with this key 112 # already exists 113 if self.has(key): 114 return False 115 116 dump = self.serializer.dumps(value) 117 record = {"id": self.key_prefix + key, "val": dump} 118 119 if timeout > 0: 120 record["expiration"] = now + datetime.timedelta(seconds=timeout) 121 self.client.update_one({"id": self.key_prefix + key}, {"$set": record}, True) 122 return True 123 124 def set(self, key: str, value: _t.Any, timeout: _t.Optional[int] = None) -> _t.Any: 125 self._expire_records() 126 return self._set(key, value, timeout=timeout, overwrite=True) 127 128 def set_many( 129 self, mapping: _t.Dict[str, _t.Any], timeout: _t.Optional[int] = None 130 ) -> _t.List[_t.Any]: 131 self._expire_records() 132 from pymongo import UpdateOne 133 134 operations = [] 135 now = self._utcnow() 136 timeout = self._normalize_timeout(timeout) 137 for key, val in mapping.items(): 138 dump = self.serializer.dumps(val) 139 140 record = {"id": self.key_prefix + key, "val": dump} 141 142 if timeout > 0: 143 record["expiration"] = now + datetime.timedelta(seconds=timeout) 144 operations.append( 145 UpdateOne({"id": self.key_prefix + key}, {"$set": record}, upsert=True), 146 ) 147 148 result = self.client.bulk_write(operations) 149 keys = list(mapping.keys()) 150 if result.bulk_api_result["nUpserted"] != len(keys): 151 query = self.client.find( 152 {"id": {"$in": [self.key_prefix + key for key in keys]}} 153 ) 154 keys = [] 155 for item in query: 156 keys.append(item["id"]) 157 return keys 158 159 def get_many(self, *keys: str) -> _t.List[_t.Any]: 160 results = self.get_dict(*keys) 161 values = [] 162 for key in keys: 163 values.append(results.get(key, None)) 164 return values 165 166 def get_dict(self, *keys: str) -> _t.Dict[str, _t.Any]: 167 self._expire_records() 168 query = self.client.find( 169 {"id": {"$in": [self.key_prefix + key for key in keys]}} 170 ) 171 results = dict.fromkeys(keys, None) 172 for item in query: 173 value = self.serializer.loads(item["val"]) 174 results[item["id"][len(self.key_prefix) :]] = value 175 return results 176 177 def add(self, key: str, value: _t.Any, timeout: _t.Optional[int] = None) -> _t.Any: 178 self._expire_records() 179 return self._set(key, value, timeout=timeout, overwrite=False) 180 181 def has(self, key: str) -> bool: 182 self._expire_records() 183 record = self.get(key) 184 return record is not None 185 186 def delete_many(self, *keys: str) -> _t.List[_t.Any]: 187 self._expire_records() 188 res = list(keys) 189 filter = {"id": {"$in": [self.key_prefix + key for key in keys]}} 190 result = self.client.delete_many(filter) 191 192 if result.deleted_count != len(keys): 193 existing_keys = [ 194 item["id"][len(self.key_prefix) :] for item in self.client.find(filter) 195 ] 196 res = [item for item in keys if item not in existing_keys] 197 198 return res 199 200 def clear(self) -> bool: 201 self.client.drop() 202 return True 203