【Python学习】操作Mongodb

#!/usr/bin/env python
# encoding: utf-8
"""
数据库操作
"""
import pymongo
from pymongo.errors import ServerSelectionTimeoutError
from mgr.common.logger import *
from mgr.common.utils import ssh_connect_with_output


def mongodb_connect(host, db_name, port=27017, user_name=None, password=None, mechanism=None, retry_count=300):
    """
    连接MongoDB数据库
    :param host: MongoDB ip地址
    :param db_name: 数据库名
    :param port: 端口
    :param user_name: 用户名
    :param password: 密码
    :param mechanism: 加密方式 'SCRAM-SHA-1' or 'MONGODB-CR'
    :return: client
    """
    client = pymongo.MongoClient(host=host, port=int(port))

    for i in range(retry_count):
        # 防止mgr重启导致数据库无法连接
        try:
            # 连接MGR数据库
            # if db_name == 'sangfor_edr':
            #     client.sangfor_edr.authenticate(user_name, password, mechanism=mechanism)
            _, ret = ssh_connect_with_output(ip=host, cmd='iptables -F;iptables -nvL', timeout=5)
            LOG_DEBUG('当前mgr iptables情况: {}'.format(ret))
            if user_name or password or mechanism:
                client[db_name].authenticate(user_name, password, mechanism=mechanism)
        except ServerSelectionTimeoutError:
            LOG_DEBUG('cannot connect mgr comgodb: {}, loop and temp to connect'.format(host))
            # 调试数据库连接不上原因
            _, ret = ssh_connect_with_output(ip=host, cmd='ps -ef|grep mongo', timeout=5)
            LOG_DEBUG('当前mgr MongoDB进程情况: {}'.format(ret))
            _, ret = ssh_connect_with_output(ip=host, cmd='/sf/edr/manager/bin/eps_services status;', timeout=5)
            LOG_DEBUG('当前mgr eps_services启动情况: {}'.format(ret))
            if i % 15 == 0 and i != 0:
                # 在ea前置有时没有更改成功，在mgr层修改一次
                ssh_connect_with_output(ip=host,
                                        cmd="sed -i 's|bind_ip = 127.0.0.1|bind_ip = 0.0.0.0|' /ac/dc/ldb/bin/mongo_config/mongo.conf"
                                            ";/sf/edr/manager/bin/eps_services restart;sleep 10;iptables -F;iptables -nvL",
                                        timeout=15)
            time.sleep(15)
        except Exception as e:
            LOG_DEBUG('auth failed', e)
            time.sleep(15)
        else:
            LOG_DEBUG('connect comgodb: {} success'.format(host))
            return client
    return False


def mongodb_close(client):
    """
    MongoDB断开连接
    :param client:MongoDB client
    :return:
    """
    try:
        client.close()
    except Exception as e:
        LOG_DEBUG(e)


def clear_collection_data(db, collection_name):
    for i in range(30):
        # 防止并发，其他节点重启mgr导致数据库连接失败
        try:
            result = db[collection_name].delete_many(filter={})
            LOG_DEBUG('清空{}集合的所有数据！影响数量：{}'.format(collection_name, result.deleted_count))
        except Exception as e:
            LOG_DEBUG('数据库操作报错原因: {}'.format(e))
            LOG_DEBUG('sleep 20s, and retry')
            time.sleep(20)
        else:
            return result


if __name__ == '__main__':
    client = mongodb_connect(host='111.111.111.111', db_name='testr', port=27017, user_name='test',
                             password='test', mechanism='SCRAM-SHA-1')
    # cur_count =client['sangfor_edr']['soft_agent_info'].find().count()
    # print cur_count
    # print client['sangfor_edr']['link_devices'].delete_many(filter={'type': 'SOC'})
    # client['sangfor_edr']['agent_info'].update_many(filter={}, update={'$set': {'patch_upgrade_state': 0}})
    # client = mongodb_connect(host='10.186.0.1', db_name='edrFuerNew', port=27017)['edrFuerNew']
    # db = client['sangfor_edr']
    # device_registration = db['device_registration']
    # ret = device_registration.find({"ip_address": {"$ne": "10.186.0.200"}}, {'_id': 1})
    # for i in ret:
    #     print i['_id']
    # ret = db['leak_info_log'].update_many(filter={}, update={"$set":{"create_time": long(1557408411)}})
    # print ret.matched_count
    # clear_collection_data(db, 'patch_agent_log')