python通过mongoengine中connect函数连接多个数据库

mongoengine支持程序同时连接多个数据库，这些数据库可以位于一个或多个mongo之中，通过alias名称区分不同的连接即可。

可以通过switch_db切换到不同的数据库，进行读写操作，switch_db其实是一个上下文管理器，通过和with语句一起使用，能够保证切换数据库的影响范围。

由于个人电脑上没有安装mongodb，以下代码示例中访问的mongodb通过mongoengine提供的mock mongodb替代，只需要在connect函数参数中加上is_mock=True参数，并且安装有mongomock package(pip install mongomock)即可。

#!/usr/bin/env python
# coding=utf-8
from mongoengine import connect
from mongoengine.document import Document
from mongoengine.fields import StringField, IntField
from mongoengine.context_managers import switch_db


# 连接不同db
connect(alias='testdbA', is_mock=True)
connect(alias='testdbB', is_mock=True)
connect(alias='testdbC', is_mock=True)


class User(Document):
    meta = {
            # db_alias用于指定User绑定的mongo连接，和connect函数中的alias对应
            'db_alias': 'testdbA',
            }
    uname = StringField(max_length=63)
    uid = StringField(max_length=64)
    age = IntField(default=0)


def create_record(uname, uid, age):
    user = User()
    user.uname = uname
    user.uid = uid
    user.age = age
    return user

def print_records(records):
    # 输出User类此时绑定的db连接名称，并输出其记录个数
    records = list(User.objects.all())
    print('{} count:{}, records:'.format(User._meta['db_alias'], len(records)))
    for rc in records:
        print('    {}|{}|{}'.format(rc['uname'], rc['uid'], rc['age']))


records = list(User.objects.all())
print_records(records)
create_record('Ace', 'AX001', 10).save()
print_records(records)

# switch到testdbB连接
with switch_db(User, 'testdbB'):
    print_records(records)
    create_record('Bob', 'BX001', 11).save()
    print_records(records)

# switch到testdbC连接
with switch_db(User, 'testdbC'):
    print_records(records)
    create_record('Carl', 'CX001', 12).save()
    print_records(records)

 

 

 运行结果如下：

$ python mongo_test.py
testdbA count:0, records:
testdbA count:1, records:
    Ace|AX001|10
testdbB count:0, records:
testdbB count:1, records:
    Bob|BX001|11
testdbC count:0, records:
testdbC count:1, records:
    Carl|CX001|12

 

switch_db其实就是一个简单的上下文管理器，源码位于mongoengine/context_managers.py当中，其实就是在switch_db对象初始化时、进入with代码块前、离开with代码块后，通过改变db_alias和collection两个字段对应的取值，实现了不同mongodb连接的切换，代码简单注释版本如下:

class switch_db(object):
    """switch_db alias context manager.

    Example ::

        # Register connections
        register_connection('default', 'mongoenginetest')
        register_connection('testdb-1', 'mongoenginetest2')

        class Group(Document):
            name = StringField()

        Group(name='test').save()  # Saves in the default db

        with switch_db(Group, 'testdb-1') as Group:
            Group(name='hello testdb!').save()  # Saves in testdb-1
    """

    def __init__(self, cls, db_alias):
        """Construct the switch_db context manager

        :param cls: the class to change the registered db
        :param db_alias: the name of the specific database to use
        """
        self.cls = cls
        # 存储切换前的collection
        self.collection = cls._get_collection()
        # 设置为要切换的db连接名称
        self.db_alias = db_alias
        # 存储切换前的原db连接名称
        self.ori_db_alias = cls._meta.get('db_alias', DEFAULT_CONNECTION_NAME)
    def __enter__(self):
        """Change the db_alias and clear the cached collection."""
        # 进入with代码块前触发
        # 修改为本次指定的db连接名称
        self.cls._meta['db_alias'] = self.db_alias
        # 将collection置为None，于是之后会重新根据新的db连接获取collection
        self.cls._collection = None
        return self.cls

    def __exit__(self, t, value, traceback):
        """Reset the db_alias and collection."""
        # 退出with代码块后触发
        # 恢复为切换前的db连接名称
        self.cls._meta['db_alias'] = self.ori_db_alias
        # 恢复为切换前的collection
        self.cls._collection = self.collection