增加Redis可选择db (#75)

proxypool/setting.py

@@ -29,11 +29,13 @@ REDIS_HOST = env.str('REDIS_HOST', '127.0.0.1')
 REDIS_PORT = env.int('REDIS_PORT', 6379)
 # redis password, if no password, set it to None
 REDIS_PASSWORD = env.str('REDIS_PASSWORD', None)
-# redis connection string, like redis://[password]@host:port or rediss://[password]@host:port
+# redis db, if no choice, set it to 0
+REDIS_DB = env.int('REDIS_DB', 1)
+# redis connection string, like redis://[password]@host:port or rediss://[password]@host:port/0
 REDIS_CONNECTION_STRING = env.str('REDIS_CONNECTION_STRING', None)
 
 if REDIS_CONNECTION_STRING:
-    REDIS_HOST, REDIS_PORT, REDIS_PASSWORD = parse_redis_connection_string(REDIS_CONNECTION_STRING)
+    REDIS_HOST, REDIS_PORT, REDIS_PASSWORD, REDIS_DB = parse_redis_connection_string(REDIS_CONNECTION_STRING)
 
 # redis hash table key name
 REDIS_KEY = env.str('REDIS_KEY', 'proxies:universal')

proxypool/storages/redis.py

@@ -1,7 +1,7 @@
 import redis
 from proxypool.exceptions import PoolEmptyException
 from proxypool.schemas.proxy import Proxy
-from proxypool.setting import REDIS_HOST, REDIS_PORT, REDIS_PASSWORD, REDIS_KEY, PROXY_SCORE_MAX, PROXY_SCORE_MIN, \
+from proxypool.setting import REDIS_HOST, REDIS_PORT, REDIS_PASSWORD, REDIS_DB, REDIS_KEY, PROXY_SCORE_MAX, PROXY_SCORE_MIN, \
     PROXY_SCORE_INIT
 from random import choice
 from typing import List
@@ -17,16 +17,16 @@ class RedisClient(object):
     """
     redis connection client of proxypool
     """
-    
-    def __init__(self, host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, **kwargs):
+
+    def __init__(self, host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_DB, **kwargs):
         """
         init redis client
         :param host: redis host
         :param port: redis port
         :param password: redis password
         """
-        self.db = redis.StrictRedis(host=host, port=port, password=password, decode_responses=True, **kwargs)
-    
+        self.db = redis.StrictRedis(host=host, port=port, password=password, db=db, decode_responses=True, **kwargs)
+
     def add(self, proxy: Proxy, score=PROXY_SCORE_INIT) -> int:
         """
         add proxy and set it to init score
@@ -41,7 +41,7 @@ class RedisClient(object):
             if IS_REDIS_VERSION_2:
                 return self.db.zadd(REDIS_KEY, score, proxy.string())
             return self.db.zadd(REDIS_KEY, {proxy.string(): score})
-    
+
     def random(self) -> Proxy:
         """
         get random proxy
@@ -60,7 +60,7 @@ class RedisClient(object):
             return convert_proxy_or_proxies(choice(proxies))
         # else raise error
         raise PoolEmptyException
-    
+
     def decrease(self, proxy: Proxy) -> int:
         """
         decrease score of proxy, if small than PROXY_SCORE_MIN, delete it
@@ -78,7 +78,7 @@ class RedisClient(object):
         else:
             logger.info(f'{proxy.string()} current score {score}, remove')
             return self.db.zrem(REDIS_KEY, proxy.string())
-    
+
     def exists(self, proxy: Proxy) -> bool:
         """
         if proxy exists
@@ -86,7 +86,7 @@ class RedisClient(object):
         :return: if exists, bool
         """
         return not self.db.zscore(REDIS_KEY, proxy.string()) is None
-    
+
     def max(self, proxy: Proxy) -> int:
         """
         set proxy to max score
@@ -97,21 +97,21 @@ class RedisClient(object):
         if IS_REDIS_VERSION_2:
             return self.db.zadd(REDIS_KEY, PROXY_SCORE_MAX, proxy.string())
         return self.db.zadd(REDIS_KEY, {proxy.string(): PROXY_SCORE_MAX})
-    
+
     def count(self) -> int:
         """
         get count of proxies
         :return: count, int
         """
         return self.db.zcard(REDIS_KEY)
-    
+
     def all(self) -> List[Proxy]:
         """
         get all proxies
         :return: list of proxies
         """
         return convert_proxy_or_proxies(self.db.zrangebyscore(REDIS_KEY, PROXY_SCORE_MIN, PROXY_SCORE_MAX))
-    
+
     def batch(self, start, end) -> List[Proxy]:
         """
         get batch of proxies

proxypool/utils/parse.py

@@ -8,6 +8,6 @@ def parse_redis_connection_string(connection_string):
     :param connection_string:
     :return:
     """
-    result = re.match('rediss?:\/\/(.*?)@(.*?):(\d+)', connection_string)
-    return result.group(2), int(result.group(3)), (result.group(1) or None) if result \
-        else ('localhost', 6379, None)
+    result = re.match('rediss?:\/\/(.*?)@(.*?):(\d+)\/(\d+)', connection_string)
+    return result.group(2), int(result.group(3)), (result.group(1) or None), (result.group(4) or 0) if result \
+        else ('localhost', 6379, None)