You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
from __future__ import annotations
|
|
import logging
|
|
from typing import TYPE_CHECKING, Optional, Dict, Tuple
|
|
from pymongo import MongoClient
|
|
from pymongo import UpdateOne
|
|
from pymongo.errors import DuplicateKeyError, BulkWriteError
|
|
|
|
if TYPE_CHECKING:
|
|
from pymongo.database import Database
|
|
from pymongo.collection import Collection
|
|
from pymongo.results import InsertManyResult, BulkWriteResult
|
|
|
|
|
|
def update_document(filter_query: dict = None, update_data: dict = None, replace: bool = True) -> Tuple[dict, dict]:
|
|
update_query = {}
|
|
if not update_data:
|
|
return {}, {}
|
|
|
|
for key, val in update_data.items():
|
|
if replace:
|
|
update_query.setdefault(
|
|
"$set", {}
|
|
).update(
|
|
{key: val}
|
|
)
|
|
else:
|
|
if isinstance(val, list):
|
|
update_query.setdefault(
|
|
"$addToSet", {}
|
|
).update({
|
|
key: {"$each": val}
|
|
})
|
|
else:
|
|
update_query.setdefault(
|
|
"$set", {}
|
|
).update(
|
|
{key: val}
|
|
)
|
|
return filter_query, update_query
|
|
|
|
|
|
class MongoDBUtils:
|
|
def __init__(self, mongo_uri, mongo_db):
|
|
self.mongo_uri = mongo_uri
|
|
self.mongo_db = mongo_db
|
|
self.client: MongoClient = None
|
|
self.db: Database = None
|
|
|
|
def _insert2db(self, items, tablename, ordered: bool = False, **kwargs) -> InsertManyResult:
|
|
collection: Collection = self.db.get_collection(tablename)
|
|
result: InsertManyResult = collection.insert_many(items, ordered=ordered, **kwargs)
|
|
return result
|
|
|
|
def _update2db(self, items, tablename, ordered: bool = False, **kwargs) -> BulkWriteResult:
|
|
collection: Collection = self.db.get_collection(tablename)
|
|
bulk_results: BulkWriteResult = collection.bulk_write(items, ordered=ordered, **kwargs)
|
|
return bulk_results
|