MongoDB - 对pymongo的简单套娃

代码

# -*- coding: utf-8 -*-

from loguru import logger
from pymongo import MongoClient


class MongoHelper:
    def __init__(self, url=None, db=None, coll=None):
        self.client = MongoClient(url) if url else MongoClient()
        self.coll = self.client.get_database(db).get_collection(coll)

    # 查找文档  可以只显示 或者不显示 某些字段
    def find(self, query: dict = None, only_show: list = None, dont_show: list = None, limit: int = None):
        query = query or {}
        only_show = only_show or []
        dont_show = dont_show or []
        show = {}
        show.update({v: 1 for v in only_show})
        show.update({v: 0 for v in dont_show})
        res = self.coll.find(query, show)
        return list(res) if limit is None else list(res.limit(limit))

    # unique存在  则不操作  反之则新增
    def insert_one(self, item: dict, unique: str = None):
        filed = unique or '_id'
        res = self.coll.update_one({filed: item[filed]}, {'$setOnInsert': item}, upsert=True)
        return True if res.raw_result.get('upserted') else False

    # unique为  _id字段  或者唯一索引字段
    def insert_many(self, items: list, unique: str = None):
        filed = unique or '_id'
        items1 = []
        temp = []
        for v in items:
            key = v[filed]
            if key not in temp:
                temp.append(key)
                items1.append(v)
        exists = [v[filed] for v in self.coll.find({filed: {'$in': temp}})]
        items2 = [v for v in items1 if v[filed] in set(temp) - set(exists)]
        # args = len(items), len(items1), len(exists), len(items2)
        # print('准备新增{}，去重后得{}，库中已有{}，最后新增{}'.format(*args))
        if items2:
            res = self.coll.insert_many(items2)
            return len(res.inserted_ids)
        else:
            return 0

    # 更新一条数据
    def update_one(self, query: dict, new: dict):
        res = self.coll.update_one(query, {'$set': new})
        return True if res.raw_result.get('nModified') else False