topic_spider/credit_china/utils/databases.py

from typing import Optional

import redis
from pymongo import MongoClient
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, coll: str, cfg: Optional[dict] = mongo_conf):
        self.client = MongoClient(host=cfg['host'], port=cfg['port'])
        self.db: Database = self.client[db]
        self.collection: Collection = self.db[coll]

    def __enter__(self):
        return self.collection

    def __exit__(self, exc_type, exc_val, exc_tb):
        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