AiOpsWechatiPad/common/redis_helper.py

import redis
import os
from config import conf
import uuid
import time
import threading
# 定义全局 redis_helper
redis_helper = None

class RedisHelper:
    def __init__(self, host='localhost', port=6379, password=None ,db=0):
        # 初始化 Redis 连接
        self.client = redis.Redis(host=host, port=port,db=db,password=password)
        self.lock_renewal_thread = None

    def set_hash(self, hash_key, data, timeout=None):
        """添加或更新哈希，并设置有效期"""
        self.client.hset(hash_key, mapping=data)
        if timeout:
            # 设置有效期（单位：秒）
            self.client.expire(hash_key, timeout)

    def get_hash(self, hash_key):
        """获取整个哈希表数据"""
        result = self.client.hgetall(hash_key)
        # 将字节数据解码成字符串格式返回
        return {k.decode('utf-8'): v.decode('utf-8') for k, v in result.items()}

    def get_hash_field(self, hash_key, field):
        """获取哈希表中的单个字段值"""
        result = self.client.hget(hash_key, field)
        return result.decode('utf-8') if result else None

    def delete_hash(self, hash_key):
        """删除整个哈希表"""
        self.client.delete(hash_key)
    
    def delete_hash_field(self, hash_key, field):
        """删除哈希表中的某个字段"""
        self.client.hdel(hash_key, field)

    def update_hash_field(self, hash_key, field, value):
        """更新哈希表中的某个字段"""
        self.client.hset(hash_key, field, value)

    def acquire_lock(self, lock_name, timeout=60):
        """
        尝试获取分布式锁，成功返回 True，失败返回 False
        :param lock_name: 锁的名称
        :param timeout: 锁的超时时间（秒）
        :return: bool
        """
        #print('获取锁')
        identifier = str(time.time())  # 使用时间戳作为唯一标识
        #if self.client.set(lock_name, identifier, nx=True, ex=timeout):
        if self.client.setnx(lock_name, identifier):
            self.client.expire(lock_name, timeout)
            self.lock_renewal_thread = threading.Thread(target=self.renew_lock, args=(lock_name, identifier, timeout))
            self.lock_renewal_thread.start()
            return True
        return False

    def renew_lock(self, lock_name, identifier, timeout):
        """
        锁的自动续期
        :param lock_name: 锁的名称
        :param identifier: 锁的唯一标识
        :param timeout: 锁的超时时间（秒）
        """
        while True:
            time.sleep(timeout / 2)
            if self.client.get(lock_name) == identifier.encode():
                self.client.expire(lock_name, timeout)
            else:
                break

    def release_lock(self, lock_name, identifier):
        """
        释放分布式锁
        :param lock_name: 锁的名称
        :param identifier: 锁的唯一标识
        """
        if self.client.get(lock_name) == identifier.encode():
            self.client.delete(lock_name)
            if self.lock_renewal_thread:
                self.lock_renewal_thread.join()


    def enqueue(self, queue_name, item):
            """
            将元素添加到队列的尾部（右侧）
            :param queue_name: 队列名称
            :param item: 要添加到队列的元素
            """
            self.client.rpush(queue_name, item)
            print(f"Enqueued: {item} to queue {queue_name}")

    def dequeue(self, queue_name):
        """
        从队列的头部（左侧）移除并返回元素
        :param queue_name: 队列名称
        :return: 移除的元素，如果队列为空则返回 None
        """
        item = self.client.lpop(queue_name)
        if item:
            print(f"Dequeued: {item.decode('utf-8')} from queue {queue_name}")
            return item.decode('utf-8')
        print(f"Queue {queue_name} is empty")
        return None

    def get_queue_length(self, queue_name):
        """
        获取队列的长度
        :param queue_name: 队列名称
        :return: 队列的长度
        """
        length = self.client.llen(queue_name)
        print(f"Queue {queue_name} length: {length}")
        return length

    def peek_queue(self, queue_name):
        """
        查看队列的头部元素，但不移除
        :param queue_name: 队列名称
        :return: 队列的头部元素，如果队列为空则返回 None
        """
        item = self.client.lrange(queue_name, 0, 0)
        if item:
            print(f"Peeked: {item[0].decode('utf-8')} from queue {queue_name}")
            return item[0].decode('utf-8')
        print(f"Queue {queue_name} is empty")
        return None

    def clear_queue(self, queue_name):
        """
        清空队列
        :param queue_name: 队列名称
        """
        self.client.delete(queue_name)
        print(f"Cleared queue {queue_name}")


def start():
    global redis_helper
    
    host=conf().get("redis_host")
    port=conf().get("redis_port")
    password=conf().get("redis_password")
    db=conf().get("redis_db")
    redis_helper = RedisHelper(host=host,port=port,password=password,db=db)