pytools/common/databases.py

import bson
import pymongo

import setting as settings

# ---------------------------------- 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 settings.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)