43 lines
1.4 KiB
Python
43 lines
1.4 KiB
Python
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})
|