|
import json
|
|
|
|
import redis
|
|
import logging
|
|
from rag import settings
|
|
from rag.utils import singleton
|
|
|
|
|
|
class Payload:
|
|
def __init__(self, consumer, queue_name, group_name, msg_id, message):
|
|
self.__consumer = consumer
|
|
self.__queue_name = queue_name
|
|
self.__group_name = group_name
|
|
self.__msg_id = msg_id
|
|
self.__message = json.loads(message['message'])
|
|
|
|
def ack(self):
|
|
try:
|
|
self.__consumer.xack(self.__queue_name, self.__group_name, self.__msg_id)
|
|
return True
|
|
except Exception as e:
|
|
logging.warning("[EXCEPTION]ack" + str(self.__queue_name) + "||" + str(e))
|
|
return False
|
|
|
|
def get_message(self):
|
|
return self.__message
|
|
|
|
|
|
@singleton
|
|
class RedisDB:
|
|
def __init__(self):
|
|
self.REDIS = None
|
|
self.config = settings.REDIS
|
|
self.__open__()
|
|
|
|
def __open__(self):
|
|
try:
|
|
self.REDIS = redis.StrictRedis(host=self.config["host"].split(":")[0],
|
|
port=int(self.config.get("host", ":6379").split(":")[1]),
|
|
db=int(self.config.get("db", 1)),
|
|
password=self.config.get("password"),
|
|
decode_responses=True)
|
|
except Exception as e:
|
|
logging.warning("Redis can't be connected.")
|
|
return self.REDIS
|
|
|
|
def health(self):
|
|
|
|
self.REDIS.ping()
|
|
a, b = 'xx', 'yy'
|
|
self.REDIS.set(a, b, 3)
|
|
|
|
if self.REDIS.get(a) == b:
|
|
return True
|
|
|
|
def is_alive(self):
|
|
return self.REDIS is not None
|
|
|
|
def exist(self, k):
|
|
if not self.REDIS: return
|
|
try:
|
|
return self.REDIS.exists(k)
|
|
except Exception as e:
|
|
logging.warning("[EXCEPTION]exist" + str(k) + "||" + str(e))
|
|
self.__open__()
|
|
|
|
def get(self, k):
|
|
if not self.REDIS: return
|
|
try:
|
|
return self.REDIS.get(k)
|
|
except Exception as e:
|
|
logging.warning("[EXCEPTION]get" + str(k) + "||" + str(e))
|
|
self.__open__()
|
|
|
|
def set_obj(self, k, obj, exp=3600):
|
|
try:
|
|
self.REDIS.set(k, json.dumps(obj, ensure_ascii=False), exp)
|
|
return True
|
|
except Exception as e:
|
|
logging.warning("[EXCEPTION]set_obj" + str(k) + "||" + str(e))
|
|
self.__open__()
|
|
return False
|
|
|
|
def set(self, k, v, exp=3600):
|
|
try:
|
|
self.REDIS.set(k, v, exp)
|
|
return True
|
|
except Exception as e:
|
|
logging.warning("[EXCEPTION]set" + str(k) + "||" + str(e))
|
|
self.__open__()
|
|
return False
|
|
|
|
def transaction(self, key, value, exp=3600):
|
|
try:
|
|
pipeline = self.REDIS.pipeline(transaction=True)
|
|
pipeline.set(key, value, exp, nx=True)
|
|
pipeline.execute()
|
|
return True
|
|
except Exception as e:
|
|
logging.warning("[EXCEPTION]set" + str(key) + "||" + str(e))
|
|
self.__open__()
|
|
return False
|
|
|
|
def queue_product(self, queue, message, exp=settings.SVR_QUEUE_RETENTION) -> bool:
|
|
for _ in range(3):
|
|
try:
|
|
payload = {"message": json.dumps(message)}
|
|
pipeline = self.REDIS.pipeline()
|
|
pipeline.xadd(queue, payload)
|
|
pipeline.expire(queue, exp)
|
|
pipeline.execute()
|
|
return True
|
|
except Exception as e:
|
|
print(e)
|
|
logging.warning("[EXCEPTION]producer" + str(queue) + "||" + str(e))
|
|
return False
|
|
|
|
def queue_consumer(self, queue_name, group_name, consumer_name, msg_id=b">") -> Payload:
|
|
try:
|
|
group_info = self.REDIS.xinfo_groups(queue_name)
|
|
if not any(e["name"] == group_name for e in group_info):
|
|
self.REDIS.xgroup_create(
|
|
queue_name,
|
|
group_name,
|
|
id="0",
|
|
mkstream=True
|
|
)
|
|
args = {
|
|
"groupname": group_name,
|
|
"consumername": consumer_name,
|
|
"count": 1,
|
|
"block": 10000,
|
|
"streams": {queue_name: msg_id},
|
|
}
|
|
messages = self.REDIS.xreadgroup(**args)
|
|
if not messages:
|
|
return None
|
|
stream, element_list = messages[0]
|
|
msg_id, payload = element_list[0]
|
|
res = Payload(self.REDIS, queue_name, group_name, msg_id, payload)
|
|
return res
|
|
except Exception as e:
|
|
if 'key' in str(e):
|
|
pass
|
|
else:
|
|
logging.warning("[EXCEPTION]consumer" + str(queue_name) + "||" + str(e))
|
|
return None
|
|
|
|
|
|
REDIS_CONN = RedisDB()
|
|
|
