match_spider/zbytb/utils/databases.py

from typing import Optional

import pymongo
import redis
from pymongo.collection import Collection
from pymongo.database import Database

from config.load import mongo_conf, redis_conf

__all__ = ['MongoDBS', 'RedisDBS']


class MongoDBS:
    """ Mongo """

    def __init__(self, db: str, collection: str, cfg: dict = mongo_conf):
        self.client = pymongo.MongoClient(host=cfg['host'], port=cfg['port'])
        self.db: Database = self.client[db]
        self.coll: Collection = self.db[collection]

    def __enter__(self):
        return self.coll

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 上下文管理器,实例调用完毕后,关闭客户端连接
        self.client.close()

    def __del__(self):
        # 实例调用完毕后,关闭客户端连接
        self.client.close()


class RedisDBS:
    """ redis """

    def __init__(self, cfg: Optional[dict] = redis_conf):
        pool = redis.ConnectionPool(
            host=cfg['host'],
            port=cfg['port'],
            password=cfg['pwd'],
            db=cfg['db']
        )
        self.__r = redis.Redis(connection_pool=pool, decode_responses=True)

    @property
    def redis(self):
        return self.__r