Powers/database/flood_db.py

from threading import RLock

from Powers.database import MongoDB

INSERTION_LOCK = RLock()


class Floods(MongoDB):
    """Class to store flood limit and action of a chat"""
    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:
            if not (curr := self.find_one({"chat_id": chat_id})):
                return self.insert_one(
                    {
                        "chat_id": chat_id,
                        "limit": limit,
                        "within": within,
                        "action": action
                    },
                )
            if (
                    limit != int(curr['limit'])
                    or within != int(curr['within'])
                    or action != str(curr['action'])
            ):
                return self.update(
                    {
                        "chat_id": chat_id
                    },
                    {
                        "limit": limit,
                        "within": within,
                        "action": action,
                    }
                )
            else:
                return

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

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

    def rm_flood(self, chat_id: int):
        with INSERTION_LOCK:
            if curr := self.find_one({"chat_id": chat_id}):
                self.delete_one({"chat_id": chat_id})
                return True
            return False