pytools/common/databases.py

import bson
import pymongo
import redis
from redis._compat import unicode, long, basestring
from redis.connection import Encoder as RedisEncoder
from redis.exceptions import DataError

# import config.load as settings
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)


# ---------------------------------- redis ----------------------------------
def redis_client(cfg=None):

    class Encoder(RedisEncoder):

        def encode(self, value):
            "Return a bytestring or bytes-like representation of the value"
            if isinstance(value, (bytes, memoryview)):
                return value
            # elif isinstance(value, bool):
            #     # special case bool since it is a subclass of int
            #     raise DataError(
            #         "Invalid input of type: 'bool'. Convert to a "
            #         "bytes, string, int or float first."
            #     )
            elif isinstance(value, float):
                value = repr(value).encode()
            elif isinstance(value, (int, long)):
                # python 2 repr() on longs is '123L', so use str() instead
                value = str(value).encode()
            elif isinstance(value, (list, dict, tuple)):
                value = unicode(value)
            elif not isinstance(value, basestring):
                # a value we don't know how to deal with. throw an error
                typename = type(value).__name__
                raise DataError(
                    "Invalid input of type: '%s'. Convert to a "
                    "bytes, string, int or float first." % typename
                )
            if isinstance(value, unicode):
                value = value.encode(self.encoding, self.encoding_errors)
            return value

    redis.connection.Encoder = Encoder

    if cfg is None:
        cfg = settings.redis_conf

    return redis.StrictRedis(host=cfg['host'], port=cfg['port'],
                             password=cfg['pwd'], db=cfg['db'],
                             decode_responses=True)