rate-sys-template/Utils/MongoUtils.py

import json
from copy import deepcopy

from bson import ObjectId
from pydantic import BaseModel
from pymongo import MongoClient
from typing import Generic, TypeVar

ModelType = TypeVar("ModelType", bound=BaseModel)


def pydantic_to_mg_dict(data: BaseModel):
    data = json.loads(data.json(exclude_none=True))
    return data


class MongoConnect:
    def __init__(self, host="localhost", port: int = 27017, db: str = None, user: str = None, password: str = None):
        self.host = host
        self.port = port
        self.db = db
        self.user = user
        self.password = password
        if self.user:
            self.client = MongoClient(host=self.host, port=self.port,
                                      username=self.user, password=self.password)
        else:
            self.client = MongoClient(host=self.host, port=self.port)
        self.db = self.client[self.db]

    def get_db(self):
        yield self.db

    def insert_data(self, dbname: str, sheet: str, data: dict):
        collection = self.client[dbname][sheet]
        item = collection.insert_one(data)
        return item.inserted_id.__str__()

    def delete_data_by_id(self, dbname: str, sheet: str, _id: str):
        collection = self.client[dbname][sheet]
        collection.delete_one({'_id': ObjectId(_id)})
        return True

    def find_data_by_id(self, dbname: str, sheet: str, _id: str):
        collection = self.client[dbname][sheet]
        return collection.find_one({'_id': ObjectId(_id)}, {"_id": False})

    def update_data_by_id(self, dbname: str, sheet: str, _id: str, data: dict):
        collection = self.client[dbname][sheet]
        collection.update_one({'_id': ObjectId(_id)}, {"$set": data})
        return True

    def replace_data_by_id(self, dbname: str, sheet: str, _id: str, data: dict):
        collection = self.client[dbname][sheet]
        collection.find_one_and_replace({'_id': ObjectId(_id)}, data)
        return True