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)