match_spider/qlm/utils/databases.py

import bson
import pymongo
import redis

from config.load import mongo_conf, redis_conf

# ---------------------------------- mongo ----------------------------------
MONGO_URI_CLIENTS = {}    # a dictionary hold all client with uri as key


def mongo_client(cfg=None, host=None, port=None, fork=False, **kwargs):
    if host is not None and port is not None:
        uri = f'mongodb://{host}:{port}'
    else:
        _cfg = (cfg or mongo_conf)
        uri = f'mongodb://{_cfg["host"]}:{_cfg["port"]}'

    if fork:
        return pymongo.MongoClient(uri, **kwargs)
    global MONGO_URI_CLIENTS
    matched_client = MONGO_URI_CLIENTS.get(uri)
    if matched_client is None:
        new_client = pymongo.MongoClient(uri, **kwargs)
        if new_client is not None:
            MONGO_URI_CLIENTS[uri] = new_client
        return new_client
    return matched_client


def mongo_database(name: str, **kw):
    client = mongo_client(**kw)
    return client.get_database(name)


def mongo_table(db: str, name: str, **kw):
    database = mongo_database(db, **kw)
    return database.get_collection(name)


def int2long(param: int):
    """int 转换成 long """
    return bson.int64.Int64(param)


def object_id(_id: str):
    return bson.objectid.ObjectId(_id)


# ---------------------------------- redis ----------------------------------
def redis_client(cfg=None):
    if cfg is None:
        cfg = redis_conf
    pool = redis.ConnectionPool(
        host=cfg['host'],
        port=cfg['port'],
        password=cfg['pwd'],
        db=cfg['db']
    )
    return redis.Redis(connection_pool=pool, decode_responses=True)