Powers/database/flood_db.py

from threading import RLock
from traceback import format_exc

from Powers import LOGGER
from Powers.database import MongoDB
from Powers.utils.msg_types import Types

INSERTION_LOCK = RLock()

class Floods(MongoDB):
    """Class to store flood limit and action of a chat"""
    try:
        db_name = "flood"

        def __init__(self):
            super().__init__(self.db_name)

        def save_flood(
            self,
            chat_id: int,
            limit: int,
            within: int,
            action: str,
        ):
            with INSERTION_LOCK:
                curr = self.find_one({"chat_id": chat_id})
                if curr:
                    if not(limit == int(curr['limit']) and within == int(curr['within']) and action == str(curr['action'])):
                        if limit != int(curr['limit']):
                            return self.update(
                                {"chat_id": chat_id,"within": within,"action": action},
                                {"limit": limit}
                            )
                        elif within != int(curr['within']):
                            return self.update(
                                {"chat_id": chat_id,"limit": limit,"action": action},
                                {"within": within}
                            )
                        elif action != str(curr['action']):
                            return self.update(
                                {"chat_id": chat_id,"limit": limit,"within": within},
                                {"action": action}
                            )
                return self.insert_one(
                    {
                        "chat_id" : chat_id,
                        "limit": limit,
                        "within": within,
                        "action" : action
                    },
                )

        def is_chat(self, chat_id: int):
            with INSERTION_LOCK:
                curr = self.find_one({"chat_id": chat_id})
                if curr:
                    action = [str(curr["limit"]), str(curr["within"]), str(curr["action"])]
                    return action
                return False

        def get_action(self, chat_id: int, limit: int, within: int, action: str):
            with INSERTION_LOCK:
                curr = self.find_one({"chat_id": chat_id, "limit": limit, "within": within, "action": action})
                if curr:
                    return curr
                return "Flood haven't set"

        def rm_flood(self, chat_id: int, limit: int, within: int, action: str):
            with INSERTION_LOCK:
                curr = self.find_one({"chat_id": chat_id, "limit": limit, "within": within, "action": action})
                if curr:
                    self.delete_one(curr)
                    return True
                return False
    except Exception as e:
        LOGGER.error(ef)
        LOGGER.error(format_exc())