AiOpsWechatiPad/common/redis_helper.py

import redis
import os
from config import conf
import uuid
import time
# 定义全局 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)

    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_key, expire_time=60, timeout=None):
            """
            获取分布式锁。
            
            Args:
                lock_key: 锁的键名。
                expire_time: 锁的有效时间（秒）。
                timeout: 最大等待时间（秒），默认为None表示无限等待。
                
            Returns:
                成功获取锁返回True，否则返回False。
            """
            identifier = str(uuid.uuid4())
            
            while True:
                # 尝试获取锁
                if self.client.setnx(lock_key, identifier):
                    self.client.expire(lock_key, expire_time)
                    return True
                
                # 检查锁是否存在且未过期
                current_value = self.client.get(lock_key)
                if not current_value:
                    continue  # 锁不存在，继续尝试获取
                
                # 检查锁是否已过期
                ttl = self.client.ttl(lock_key)
                if ttl == -2:  # 锁已过期
                    continue
                
                # 如果超时时间设置且已经超出，则返回False
                if timeout is not None:
                    start_time = time.time()
                    while (time.time() - start_time) < timeout:
                        time.sleep(0.1)
                        return self.acquire_lock(lock_key, expire_time, timeout)
                    else:
                        # 超时，放弃获取锁
                        return False
                
                # 等待一段时间后重试
                time.sleep(0.1)

    def release_lock(self, lock_key):
        """
        释放分布式锁。
        
        Args:
            lock_key: 锁的键名。
            
        Returns:
            成功释放返回True，否则返回False。
        """
        script = """
        if redis.call("get", KEYS[1]) == ARGV[1] then
            return redis.call("del", KEYS[1])
        else
            return 0
        end
        """
        current_value = self.client.get(lock_key)
        if not current_value:
            return False
        
        identifier = str(uuid.uuid4())  # 这里应替换为获取锁时使用的标识符
        result = self.client.eval(script, [lock_key], [identifier])
        return result == 1
        

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)