MongoDB Python官方驱动 PyMongo 的简单封装

最近，需要使用 python 对 mongodb 做一些简单的操作，不想使用各种繁重的框架。出于可重用性的考虑，想对 mongodb python 官方驱动 pymongo 做下简单封装，百度一如既往的未能给我一个满意的结果，于是有了下文。

---

【正文】

pymongo，mongodb python官方驱动

• docs: https://api.mongodb.com/python/current/index.html
• github: https://github.com/mongodb/mongo-python-driver

pymongo 驱动几乎支持 mongodb 的全部特性，可以连接单个的 mongodb 数据库、副本集和分片集群。从提供的api角度来看，pymongo package是其核心，包含对数据库的各种操作。本文将介绍一个简单封装类 dbmanager。主要特性：对数据库和集合的操作确保其存在性；支持pymongo的原生操作，包括基本的crud操作、批量操作、mapreduce、多线程和多进程等；支持因果一致性会话和事务的流水线操作，并给出简单示例。

mongoclient

mongo_client 提供了连接 mongodb 的mongoclient类：
class pymongo.mongo_client.mongoclient(host='localhost', port=27017, document_class=dict, tz_aware=false, connect=true, **kwargs)
每个 mongoclient 实例 client （下文简称 client）都维护一个内建的连接池，默认 maxpoolsize 大小100。对于多线程的操作，连接池会给每一个线程一个 socket 连接，直到达到最大的连接数，后续的线程会阻塞以等待有可用的连接被释放。client 对 mongodb 拓扑结构中的每个server 还维护一个额外的连接来监听 server 的状态。

下面的 new_mongo_client函数用于获取一个数据库连接的 client。其中，client.admin.command('ismaster') 用来检查 server 的可用状态，简单省事不需要认证。

def new_mongo_client(uri, **kwargs):
    """create new pymongo.mongo_client.mongoclient instance. do not use it directly."""

    try:
        client = mongoclient(uri, maxpoolsize=1024, **kwargs)
        client.admin.command('ismaster')  # the ismaster command is cheap and does not require auth.
    except connectionfailure:
        logging.error("new_mongo_client(): server not available, please check you uri: {}".format(uri))
        return none
    else:
        return client

pymongo 不是进程（fork-safe）安全的，但在一个进程中是线程安全（thread-safe）的。因此常见的场景是，对于一个mongodb 环境，为每一个进程中创建一个 client ，后面所有的数据库操作都使用这一个实例，包括多线程操作。永远不要为每一次操作都创建一个 mongoclient 实例，使用完后调用 mongoclient.close() 方法，这样没有必要而且会非常浪费性能。

鉴于以上原因，一般不宜直接使用new_mongo_client函数获取 client，而是进一步封装为get_mongo_client方法。 其中全局常量 uri_client_dict 保持着数据库 uri 字符串与对应 clinet 的字典，一个 uri 对应一个 client 。代码如下：

mongo_uri_default = 'mongodb://localhost:27017/admin'
uri_client_dict = {}    # a dictionary hold all client with uri as key
def get_mongo_client(uri=mongo_uri_default, fork=false, **kwargs):
    """get pymongo.mongo_client.mongoclient instance. one mongodb uri, one client.

    @:param uri: mongodb uri
    @:param fork: for fork-safe in multiprocess case, if fork=true, return a new mongoclient instance, default false.
    @:param kwargs: refer to pymongo.mongo_client.mongoclient kwargs
    """
    if fork:
        return new_mongo_client(uri, **kwargs)
    global uri_client_dict
    matched_client = uri_client_dict.get(uri)
    if matched_client is none:  # no matched client
        new_client = new_mongo_client(uri, **kwargs)
        if new_client is not none:
            uri_client_dict[uri] = new_client
        return new_client
    return matched_client

确保 database 和 collection 的存在

pymongo 有个特性：对于不存在的数据库、集合上的查询不会报错。如下，ipython中演示在不存在xxdb 数据库和 xxcollection 集合上的操作：

in [1]: from pymongo import mongoclient
in [2]: client = mongoclient()          # default uri is 'mongodb://localhost:27017/admin'
in [3]: db = client.get_database('xxdb')    # database(mongoclient(host=['localhost:27017'], document_class=dict, tz_aware=false, connect=true), u'xxdb')
in [4]: coll = db.get_collection('xxcollection')  #  collection(database(mongoclient(host=['localhost:27017'], document_class=dict, tz_aware=false, connect=true), u'xxdb'), u'xxcollection')
in [5]: coll.find_one()                 # note: no tip, no error, no exception， return none
in [6]: coll.insert_one({'hello' : 'what a fucking feature'})
out[6]: <pymongo.results.insertoneresult at 0x524ccc8>
in [7]: coll.find_one()
out[7]: {u'_id': objectid('5c31c807bb048515b814d719'), u'hello': u'what a fucking feature'}

这对于手误写错数据库或集合名字后进行的后续操作，简直就是灾难。鉴于此因，有必要对获取数据库或集合时加上确认保护。
下面对于获取数据库，使用 mongoclient.list_database_names() 获取所有的数据库名字，如果数据库名称不在其中，则返回none。同样的道理，对于集合使用 database.list_collection_names()。注：由于用户权限问题造成的获取数据库或集合列表的操作报错的情况，默认不加确认保护。

def get_existing_db(client, db_name):
    """get existing pymongo.database.database instance.

    @:param client: pymongo.mongo_client.mongoclient instance
    @:param db_name: database name wanted
    """

    if client is none:
        logging.error('client {} is none'.format(client))
        return none
    try:
        db_available_list = client.list_database_names()
    except pymongoerror as e:
        logging.error('client: {}, db_name: {}, client.list_database_names() error: {}'.
                      format(client, db_name, repr(e)))
    else:
        if db_name not in db_available_list:
            logging.error('client {} has no db named {}'.format(client, db_name))
            return none
    db = client.get_database(db_name)
    return db

def get_existing_coll(db, coll_name):
    """get existing pymongo.collection.collection instance.

    @:param client: pymongo.mongo_client.mongoclient instance
    @:param coll_name: collection name wanted
    """

    if db is none:
        logging.error('db {} is none'.format(db))
        return none
    try:
        coll_available_list = db.list_collection_names()
    except pymongoerror as e:
        logging.error('db: {}, coll_name: {}, db.list_collection_names() error: {}'.
                      format(db, coll_name, repr(e)))
    else:
        if coll_name not in coll_available_list:
            logging.error('db {} has no collection named {}'.format(db, coll_name))
            return none
    coll = db.get_collection(coll_name)
    return coll

pymongo 封装类 dbmanger

前文的冗长铺垫主要是为了引入这个 pymongo 驱动封装类 dbmanger。

dbmanger 类的实例保持的状态有mongoclient实例 self.client, 数据库self.db 和 集合self.coll，并通过属性（property）对外开放。pymongo 原生的方法对这里的 client, db 和 coll 同样适用。client 由类的构造器调用上文的get_mongo_client方法获取，db 和 coll 即可通过类的构造器获取也可通过 self.db_name 和 self.coll_name 这些 setter 来切换。
dbmanger 类的实例持有的方法 self.create_coll(self, db_name, coll_name), session_pipeline(self, pipeline) 和 transaction_pipeline(self, pipeline)。后两种方法在下一节再具体解释。

class dbmanager:
    """a safe and simple pymongo packaging class ensuring existing database and collection.

    operations:
    mongoclient level operations: https://api.mongodb.com/python/current/api/pymongo/mongo_client.html
    database level operations: https://api.mongodb.com/python/current/api/pymongo/database.html
    collection level operations: https://api.mongodb.com/python/current/api/pymongo/collection.html
    """
    __default_uri = 'mongodb://localhost:27017/admin'
    __default_db_name = 'test'
    __default_coll_name = 'test'

    def __init__(self, uri=__default_uri, db_name=__default_db_name, coll_name=__default_coll_name, **kwargs):
        self.__uri = uri
        self.__db_name = db_name
        self.__coll_name = coll_name
        self.__client = get_mongo_client(uri, **kwargs)
        self.__db = get_existing_db(self.__client, db_name)
        self.__coll = get_existing_coll(self.__db, coll_name)

    def __str__(self):
        return u'uri: {}, db_name: {}, coll_name: {}, id_client: {}, client: {}, db: {}, coll: {}'.format(
            self.uri, self.db_name, self.coll_name, id(self.client), self.client, self.db, self.coll)

    @property
    def uri(self):
        return self.__uri

    @property
    def db_name(self):
        return self.__db_name

    @property
    def coll_name(self):
        return self.__coll_name

    @db_name.setter
    def db_name(self, db_name):
        self.__db_name = db_name
        self.__db = get_existing_db(self.__client, db_name)

    @coll_name.setter
    def coll_name(self, coll_name):
        self.__coll_name = coll_name
        self.__coll = get_existing_coll(self.__db, coll_name)

    @property
    def client(self):
        return self.__client

    @property
    def db(self):
        return self.__db

    @property
    def coll(self):
        # always use the current instance self.__db
        self.__coll = get_existing_coll(self.__db, self.__coll_name)
        return self.__coll

    def create_coll(self, db_name, coll_name):
        """create new collection with new or existing database"""
        if self.__client is none:
            return none
        try:
            return self.__client.get_database(db_name).create_collection(coll_name)
        except collectioninvalid:
            logging.error('collection {} already exists in database {}'.format(coll_name, db_name))
            return none

    def session_pipeline(self, pipeline):
        if self.__client is none:
            logging.error('client is none in session_pipeline: {}'.format(self.__client))
            return none
        with self.__client.start_session(causal_consistency=true) as session:
            result = []
            for operation in pipeline:
                try:
                    if operation.level == 'client':
                        target = self.__client
                    elif operation.level == 'db':
                        target = self.__db
                    elif operation.level == 'coll':
                        target = self.__coll

                    operation_name = operation.operation_name
                    args = operation.args
                    kwargs = operation.kwargs
                    operator = getattr(target, operation_name)
                    if type(args) == tuple:
                        ops_rst = operator(*args, session=session, **kwargs)
                    else:
                        ops_rst = operator(args, session=session, **kwargs)

                    if operation.callback is not none:
                        operation.out = operation.callback(ops_rst)
                    else:
                        operation.out = ops_rst

                except exception as e:
                    logging.error('{} {} exception, session_pipeline args: {}, kwargs: {}'.format(
                        target, operation, args, kwargs))
                    logging.error('session_pipeline exception: {}'.format(repr(e)))
                result.append(operation)
            return result

    # https://api.mongodb.com/python/current/api/pymongo/client_session.html#transactions
    def transaction_pipeline(self, pipeline):
        if self.__client is none:
            logging.error('client is none in transaction_pipeline: {}'.format(self.__client))
            return none
        with self.__client.start_session(causal_consistency=true) as session:
            with session.start_transaction():
                result = []
                for operation in pipeline:
                    try:
                        if operation.level == 'client':
                            target = self.__client
                        elif operation.level == 'db':
                            target = self.__db
                        elif operation.level == 'coll':
                            target = self.__coll
                        operation_name = operation.operation_name
                        args = operation.args
                        kwargs = operation.kwargs
                        operator = getattr(target, operation_name)
                        if type(args) == tuple:
                            ops_rst = operator(*args, session=session, **kwargs)
                        else:
                            ops_rst = operator(args, session=session, **kwargs)

                        if operation.callback is not none:
                            operation.out = operation.callback(ops_rst)
                        else:
                            operation.out = ops_rst
                    except exception as e:
                        logging.error('{} {} exception, transaction_pipeline args: {}, kwargs: {}'.format(
                            target, operation, args, kwargs))
                        logging.error('transaction_pipeline exception: {}'.format(repr(e)))
                        raise exception(repr(e))
                    result.append(operation)
                return result

这里给出一些例子来说明 dbmanager的使用方法。

• 创建集合、切换数据库或集合：

# get dbmanger instance
var dbm = dbmanager('mongodb://localhost:27017/admin') # db_name, coll_name default 'test'
dbm.create_coll('testdb', 'testcollection')
# change db or coll
dbm.db_name = 'testdb' # dbm.db (test -> testdb) and dbm.coll (test.testcollection-> testdb.testcollection) will be changed at the same time
dbm.coll_nmae = 'testcollection' # dbm.coll (test.test-> test.testcollection) will be change at the same time

• 基本的操作，crud:

# simple manipulation operation
dbm.coll.insert_one({'hello': 'world'})
print(dbm.coll.find_one())   # {'_id': objectid('...'), 'hello': 'world'}
dbm.coll.update_one({'hello': 'world'}, {'hello': 'hell'})

# bulk operation
from pymongo import insertone, deleteone, replaceone, replaceone
dbm.coll.bulk_write([insertone({'y':1}), deleteone({'x':1}), replaceone({{'w':1}, {'z':1}, upsert=true})])

# simple managing operation
import pymongo
dbm.coll.create_index([('hello', pymongo.descending)], background=true)
dbm.client.list_database_names()
dbm.db.list_collection_names()

• 线程并发，进程并行：

# thread concurrent
import threading
def fun(uri, db_name, coll_name):
    # new dbmanager instance avoid variable competition
    dbm = dbmanager(uri, db_name, coll_name)
    pass
t = threading.thread(target=func, args=(uri, db_name, coll_name))
t.start()

# multiprocess parallel
import multiprocessing
def func(uri, db_name, coll_name):
    # new process, new client with fork=true parameter, and new dbmanager instance.
    dbm = dbmanager(uri, db_name, coll_name, fork=true)
    # do something with db.
    pass
proc = multiprocessing.process(target=func, args=(uri, db_name, coll_name))
proc.start()

• mapreduce ：

# mapreduce
from bson.code import code
mapper = code('''
function () {...}
''')
reducer = code('''
function (key, value) {...}
''')
rst = dbm.coll.inline_map_reduce(mapper, reducer)

对 mongodb 一致性会话（session）和 事务（transaction）的支持

mongodb reference

• docs: https://docs.mongodb.com/manual/
• causal-consistency session: https://docs.mongodb.com/manual/core/read-isolation-consistency-recency/#causal-consistency
• transation: https://docs.mongodb.com/manual/core/transactions/#transactions

会话（session），是对数据库连接的一种逻辑表示。从mongodb 3.6开始，mongodb引入了客户端会话（client session），并在其中加入了对操作的因果一致性(causal-consistency)的支持。因此，更准确地说，这里 dbmanger 类封装的其实是因果一致性的会话，即client.start_session(causal_consistency=true)。不过，一致性能够保证的前提是客户端的应用应保证在一个会话中只有一个线程（thread）在做这些操作。在一个客户端会话中，多个顺序的读写操作得到的结果与它们的执行顺序将是因果一致的，读写的设置都自动设为 "majority"。应用场景：先写后读，先读后写，一致性的写，一致性的读（read your writes，writes follow reads，monotonic writes， monotonic reads）。客户端会话与服务端会话（server session）进行交互。从3.6版本开始，mongodb驱动将所有的操作都关联到服务端会话。服务端会话是客户端会话顺序操作因果一致性和重试写操作的得以支持的底层框架。

mongodb 对单个文档的操作时是原子性的（atomic）。原子性是指一个操作的结果要么有要么没有，不可再切割，换句话说叫 “all or nothing”。从mongodb 4.0开始，副本集（replica set）开始支持多个文档级别的原子性，即多文档事务（muti-document transaction）。在同一个事务中，对跨越不同数据库或集合下的多个文档操作，如果全部操作成功，则该事务被成功提交（commit）；如果某些操作出现失败，则整个事务会终止（abort），操作中对数据库的改动会被丢弃。只有在事务被成功提交之后，操作的结果才能被事务外看到，事务正在进行或者事务失败，其中的操作对外都不可见。单个mongod服务和分片集群（sharded cluster）暂不支持事务。mongodb官方预计在4.2版本左右对分片集群加入对事务的支持。另外，需要注意的是，多文档事务会引入更大的性能开销，在场景允许的情况下，尽可能考虑用嵌套文档或数组的单文档操作方式来解决问题。

会话和事务的主要应用场景其实都是多个的时序性操作，即流水线形式。因此 dbmanager 加入了session_pipeline(self, pipeline) 和 transaction_pipeline(self, pipeline)的操作方法。首先引入表征操作的类operation，描述一个操作作用的层次（client, db或coll）、操作方法、参数和操作结果需要调用的回调函数，见名知意，不再赘解。多个操作 operation 类的实例构成的list 为pipeline, 作为session_pipeline(self, pipeline) 和 transaction_pipeline(self, pipeline) 的输入参数。pipeline 操作的每一步的输出会写入到对应operation 类的实例的out属性中。

class operation:
    """operation for constructing sequential pipeline. only used in dbmanager.session_pipeline() or transaction_pipeline().

    constructor parameters:
    level: <'client' | 'db' | 'coll'> indicating different operation level, mongoclient, database, collection
    operation_name: literally, the name of operation on specific level
    args: position arguments the operation need. require the first parameter or a tuple of parameters of the operation.
    kwargs: key word arguments the operation need.
    callback: callback function for operation result
    
    examples:
    # pymongo.collection.collection.find(filter, projection, skip=none, limit=none,...)
    operation('coll', 'find', {'x': 5}) only filter parameter, equivalent to:
    operation('coll', 'find', args={'x': 5}) or operation('coll', 'find', kwargs={filter: {'x': 5}})

    operation('coll', 'find', ({'x': 5},{'_id': 0}) {'limit':100}), equivalent to:
    operation('coll', 'find', args=({'x': 5},{'_id': 0}, none, {'limit':100}) ), or
    operation('coll', 'find', kwargs={'filter':{'x': 5}, 'projection': {'_id': 0},'limit':100})

    def cursor_callback(cursor):
        return cursor.distinct('hello')
    operation('coll', 'find', kwargs={'limit': 2}, callback=cursor_callback)
    """
    def __init__(self, level, operation_name, args=(), kwargs={}, callback=none):
        self.level = level
        self.operation_name = operation_name
        self.args = args
        if kwargs is none:
            self.kwargs = none
        else:
            self.kwargs = kwargs
        self.callback = callback
        self.out = none

基于 dbmanager 和 operation 的因果一致性的会话和事务的简单示例如下：

# causal-consistency session or transaction pipeline operation
def cursor_callback(cursor):
    return cursor.distinct('hello')
op_1 = operation('coll', 'insert_one', {'hello': 'heaven'})
op_2 = operation('coll', 'insert_one', {'hello': 'hell'})
op_3 = operation('coll', 'insert_one', {'hello': 'god'})
op_4 = operation('coll', 'find', kwargs={'limit': 2}, callback=cursor_callback)
op_5 = operation('coll', 'find_one', {'hello': 'god'})
pipeline = [op_1, op_2, op_3, op_4, op_5]
ops = dbm.transaction_pipeline(pipeline) # only on replica set deployment
# ops = dbm.session_pipeline(pipeline) # can be standalone, replica set or sharded cluster.
for op in ops:
    print(op.out)

【正文完】

注：内容同步自同名csdn博客：https://blog.csdn.net/fzlulee/article/details/85944967