SaltStack源碼分析之使用MongoDB模塊

MongoDB模塊/usr/lib/python2.6/site-packages/salt/modules/mongodb.py

MongoDB模塊會先去檢查是否安裝有PyMongo模塊

# -*- coding: utf-8 -*-
'''
Module to provide MongoDB functionality to Salt

:configuration: This module uses PyMongo, and accepts configuration details as
    parameters as well as configuration settings::

        mongodb.host: 'localhost'
        mongodb.port: 27017
        mongodb.user: ''
        mongodb.password: ''

    This data can also be passed into pillar. Options passed into opts will
    overwrite options passed into pillar.
'''

# Import python libs
import logging

# Import salt libs
from salt._compat import string_types

# Import third party libs
try:
    import pymongo
    HAS_MONGODB = True
except ImportError:
    HAS_MONGODB = False

log = logging.getLogger(__name__)

引入logging模塊處理日志，再引入salt._compat.string_types

對檢查pymongo進行異常處理，如果導入pymongo成功，設置HAS_MONGODB變量為True，導入失敗設置為False

def __virtual__():
    '''
    Only load this module if pymongo is installed
    '''
    if HAS_MONGODB:
        return 'mongodb'
    else:
        return False

如果定義一個__virtual__函數，可以控制是否允許這個模塊可見。如果返回False，SaltStack將會忽略這個模塊，如果返回一個字符串，SaltStack將使用這個字符串。

def _connect(user=None, password=None, host=None, port=None, database='admin'):
    '''
    Returns a tuple of (user, host, port) with config, pillar, or default
    values assigned to missing values.
    '''
    if not user:
        user = __salt__['config.option']('mongodb.user')
    if not password:
        password = __salt__['config.option']('mongodb.password')
    if not host:
        host = __salt__['config.option']('mongodb.host')
    if not port:
        port = __salt__['config.option']('mongodb.port')

    try:
        conn = pymongo.connection.Connection(host=host, port=port)
        mdb = pymongo.database.Database(conn, database)
        if user and password:
            mdb.authenticate(user, password)
    except pymongo.errors.PyMongoError:
        log.error('Error connecting to database {0}'.format(database))
        return False

    return conn

pymongo.connection.Connection

這個方法PyMongo已經棄用，使用新的MongoClient代替。

def db_list(user=None, password=None, host=None, port=None):
    '''
    List all Mongodb databases

    CLI Example:

    .. code-block:: bash

        salt '*' mongodb.db_list <user> <password> <host> <port>
    '''
    conn = _connect(user, password, host, port)
    if not conn:
        return 'Failed to connect to mongo database'

    try:
        log.info('Listing databases')
        return conn.database_names()
    except pymongo.errors.PyMongoError as err:
        log.error(err)
        return str(err)

列出所有的MongoDB數據庫

$ sudo salt '*' mongodb.db_list '' '' 10.10.41.17 28018
localhost.localdomain:
    'mongodb.db_list' is not available.
gintama-qa-server:
    - pay
    - clan
    - test
    - game1
    - admin
    - gam_server_1

SaltStack會先去檢測是否包含有PyMongo模塊，如果有就執行沒有就不執行。

def db_exists(name, user=None, password=None, host=None, port=None):
    '''
    Checks if a database exists in Mongodb

    CLI Example:

    .. code-block:: bash

        salt '*' mongodb.db_exists <name> <user> <password> <host> <port>
    '''
    dbs = db_list(user, password, host, port)

    if isinstance(dbs, string_types):
        return False

    return name in dbs

檢查一個庫是否存在

$ sudo salt 'gintama-qa-server' mongodb.db_exists pay  '' '' 10.10.41.17 28018
gintama-qa-server:
    True
$ sudo salt 'gintama-qa-server' mongodb.db_exists pay3  '' '' 10.10.41.17 28018
gintama-qa-server:
    False

def db_remove(name, user=None, password=None, host=None, port=None):
    '''
    Remove a Mongodb database

    CLI Example:

    .. code-block:: bash

        salt '*' mongodb.db_remove <name> <user> <password> <host> <port>
    '''
    conn = _connect(user, password, host, port)
    if not conn:
        return 'Failed to connect to mongo database'

    try:
        log.info('Removing database {0}'.format(name))
        conn.drop_database(name)
    except pymongo.errors.PyMongoError as err:
        log.error(
            'Removing database {0} failed with error: {1}'.format(
                name, str(err)
            )
        )
        return str(err)

    return True

刪除一個庫

def user_list(user=None, password=None, host=None, port=None, database='admin'):
    '''
    List users of a Mongodb database

    CLI Example:

    .. code-block:: bash

        salt '*' mongodb.user_list <name> <user> <password> <host> <port> <database>
    '''
    conn = _connect(user, password, host, port)
    if not conn:
        return 'Failed to connect to mongo database'

    try:
        log.info('Listing users')
        mdb = pymongo.database.Database(conn, database)

        output = []

        for user in mdb.system.users.find():
            output.append([
                ('user', user['user']),
                ('readOnly', user.get('readOnly', 'None'))
            ])
        return output

    except pymongo.errors.PyMongoError as err:
        log.error(
            'Listing users failed with error: {0}'.format(
                str(err)
            )
        )
        return str(err)

列出賬號

def user_exists(name, user=None, password=None, host=None, port=None,
                database='admin'):
    '''
    Checks if a user exists in Mongodb

    CLI Example:

    .. code-block:: bash

        salt '*' mongodb.user_exists <name> <user> <password> <host> <port> <database>
    '''
    users = user_list(user, password, host, port, database)
    for user in users:
        if name == dict(user).get('user'):
            return True

    return False

查看賬號是否存在

def user_create(name, passwd, user=None, password=None, host=None, port=None,
                database='admin'):
    '''
    Create a Mongodb user

    CLI Example:

    .. code-block:: bash

        salt '*' mongodb.user_create <name> <user> <password> <host> <port> <database>
    '''
    conn = _connect(user, password, host, port)
    if not conn:
        return 'Failed to connect to mongo database'

    try:
        log.info('Creating user {0}'.format(name))
        mdb = pymongo.database.Database(conn, database)
        mdb.add_user(name, passwd)
    except pymongo.errors.PyMongoError as err:
        log.error(
            'Creating database {0} failed with error: {1}'.format(
                name, str(err)
            )
        )
        return str(err)
    return True

創建賬號

def user_remove(name, user=None, password=None, host=None, port=None,
                database='admin'):
    '''
    Remove a Mongodb user

    CLI Example:

    .. code-block:: bash

        salt '*' mongodb.user_remove <name> <user> <password> <host> <port> <database>
    '''
    conn = _connect(user, password, host, port)
    if not conn:
        return 'Failed to connect to mongo database'

    try:
        log.info('Removing user {0}'.format(name))
        mdb = pymongo.database.Database(conn, database)
        mdb.remove_user(name)
    except pymongo.errors.PyMongoError as err:
        log.error(
            'Creating database {0} failed with error: {1}'.format(
                name, str(err)
            )
        )
        return str(err)

    return True

測試情況

$ sudo salt 'gintama-qa-server' mongodb.user_create gintama gintama123 '' '' 10.10.41.17 28018 taiwan_game1
gintama-qa-server:
    True
$ sudo salt 'gintama-qa-server' mongodb.user_exists gintama  '' '' 10.10.41.17 28018 taiwan_game1
gintama-qa-server:
    True
$ sudo salt 'gintama-qa-server' mongodb.user_list   '' '' 10.10.41.17 28018 taiwan_game1
gintama-qa-server:
    |_
      |_
        - user
        - gintama
      |_
        - readOnly
        - None

$ sudo salt 'gintama-qa-server' mongodb.user_remove gintama  '' '' 10.10.41.17 28018 taiwan_game1
gintama-qa-server:
    True
$ sudo salt 'gintama-qa-server' mongodb.user_list   '' '' 10.10.41.17 28018 taiwan_game1
gintama-qa-server:
$ sudo salt 'gintama-qa-server' mongodb.user_exists gintama  '' '' 10.10.41.17 28018 taiwan_game1
gintama-qa-server:
    False