database.py

import os
from os import getenv

from dotenv import load_dotenv

import datetime
import time

from motor import motor_asyncio
from motor.core import AgnosticClient

from logger import LOGS

load_dotenv()
MONGO_URL = os.environ["MONGO_URL"]

class Database:
    def __init__(self, uri: str) -> None:
        self.client: AgnosticClient = motor_asyncio.AsyncIOMotorClient(uri)
        self.db = self.client["Akeno"]
        self.backup_chatbot = self.db["metai"]
    
    async def connect(self):
        try:
            await self.client.admin.command("ping")
            LOGS.info(f"Database Connection Established!")
        except Exception as e:
            LOGS.info(f"DatabaseErr: {e} ")
            quit(1)

    async def _close(self):
        await self.client.close()

    def get_datetime(self) -> str:
        return datetime.datetime.now().strftime("%d/%m/%Y - %H:%M")

    async def set_env(self, name: str, value: str) -> None:
        await self.env.update_one(
            {"name": name}, {"$set": {"value": value}}, upsert=True
        )

    async def get_env(self, name: str) -> str | None:
        if await self.is_env(name):
            data = await self.env.find_one({"name": name})
            return data["value"]
        return None

    async def rm_env(self, name: str) -> None:
        await self.env.delete_one({"name": name})

    async def is_env(self, name: str) -> bool:
        if await self.env.find_one({"name": name}):
            return True
        return False

    async def get_all_env(self) -> list:
        return [i async for i in self.env.find({})]

    async def _update_openai_chat_in_db(self, user_id, chatbot_chat):
        await self.backup_chatbot.update_one(
            {"user_id": user_id},
            {"$set": {"chatbot_chat": chatbot_chat}},
            upsert=True
        )

    async def _get_openai_chat_from_db(self, user_id):
        user_data = await self.backup_chatbot.find_one({"user_id": user_id})
        return user_data.get("chatbot_chat", []) if user_data else []

    async def _clear_chatbot_history_in_db(self, user_id):
        unset_clear = {"chatbot_chat": None}
        return await self.backup_chatbot.update_one({"user_id": user_id}, {"$unset": unset_clear})

    async def _clear_chatbot_database(self, user_id):
        result = await self._clear_chatbot_history_in_db(user_id)
        if result.modified_count > 0:
            return "Chat history cleared successfully."
        else:
            return "No chat history found to clear."

db = Database(MONGO_URL)