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