sysom1/environment/1_sdk/cmg_redis/utils.py

94 lines
2.3 KiB
Python

# -*- coding: utf-8 -*- #
"""
Time 2022/7/26 17:04
Author: mingfeng (SunnyQjm)
Email mfeng@linux.alibaba.com
File utils.py
Description:
"""
import uuid
import time
from redis import Redis
from redis.exceptions import ConnectionError
from cmg_base import CmgUrl, CmgException
def do_connect(url: str) -> Redis:
cec_url = CmgUrl.parse(url)
return do_connect_by_cmg_url(cec_url)
def do_connect_by_cmg_url(cec_url: CmgUrl) -> Redis:
host_port = cec_url.netloc.split(":")
if len(host_port) != 2:
raise CmgException(
f"Not valid host:port => {host_port[0]}:{host_port[1]}")
host, port = host_port[0], int(host_port[1])
try:
redis_client = Redis(host=host, port=port, db=0, decode_responses=True,
**cec_url.params)
except ConnectionError as e:
raise CmgException(e)
return redis_client
class RedisLocker:
"""
This is a simple redis lock implement
Args:
redis_client(Redis): Redis client
key(str): Locker key
ex(int): Expire time, seconds
"""
def __init__(self, redis_client: Redis, key: str, ex: int = 10):
self._redis_client = redis_client
self._key = key
self._ex = ex
self._get_locker = False
self._lock_id = self.generate_consumer_id()
def __enter__(self):
return self.lock()
def __exit__(self, exc_type, exc_val, exc_tb):
return self.unlock()
@classmethod
def generate_consumer_id(cls) -> str:
"""Generate one random lock ID
Generate a random lock ID
Returns:
str: The generated consumer ID
Examples:
>>> RedisLocker.generate_consumer_id()
30e2fda7-d4b2-48b0-9338-78ff389648e7
"""
return str(uuid.uuid4())
def unlock(self):
"""Unlock
Returns:
"""
if self._redis_client.get(self._key) == self._lock_id:
return self._redis_client.delete(self._key) == 1
return False
def lock(self) -> bool:
"""Lock
Returns:
"""
while self._redis_client.set(
self._key, self._lock_id, nx=True, ex=self._ex
) != 1:
time.sleep(0.1)
return True