grpc

utils/storage.py

@@ -0,0 +1,42 @@
+from cachetools import TTLCache
+import os
+
+from utils.mongo import mongo
+
+CACHE_SIZE = int(os.getenv("CACHE_SIZE", 1000))
+CACHE_TTL = int(os.getenv("CACHE_TTL", 3600))
+
+cache = TTLCache(CACHE_SIZE, CACHE_TTL)
+
+
+def get_chat_info(chat_id: int) -> dict:
+    cached_info = cache.get(chat_id)
+    if cached_info is not None:
+        return cached_info
+    mongo_info = mongo.chats_collection.find_one({"chat_id": chat_id})
+    if mongo_info is not None:
+        cache[chat_id] = mongo_info
+        return mongo_info
+    chat_info = {"chat_id": chat_id, "state": "default", "probability": 100}
+    mongo.chats_collection.insert_one(chat_info)
+    cache[chat_id] = chat_info
+    return chat_info
+
+
+def set_values(chat_id: int, **values):
+    cached_info = cache.get(chat_id)
+    if cached_info is None:
+        mongo_info = mongo.chats_collection.find_one({"chat_id": chat_id})
+        if mongo_info is None:
+            chat_info = {"chat_id": chat_id, "state": "default", "probability": 100}
+            chat_info.update(values)
+            mongo.chats_collection.insert_one(chat_info)
+            cache[chat_id] = chat_info
+        else:
+            mongo.chats_collection.update_one({"chat_id": chat_id}, {"$set": values})
+            mongo_info = dict(mongo_info)
+            mongo_info.update(values)
+            cache[chat_id] = mongo_info
+    else:
+        cached_info.update(values)
+        mongo.chats_collection.update_one({"chat_id": chat_id}, {"$set": values})