57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
|
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
|