uncloud/ucloud/common/etcd_wrapper.py

97 lines
3 KiB
Python

import etcd3
import json
import queue
import copy
from collections import namedtuple
from functools import wraps
from . import logger
PseudoEtcdMeta = namedtuple('PseudoEtcdMeta', ['key'])
class EtcdEntry:
# key: str
# value: str
def __init__(self, meta, value, value_in_json=False):
self.key = meta.key.decode('utf-8')
self.value = value.decode('utf-8')
if value_in_json:
self.value = json.loads(self.value)
def readable_errors(func):
@wraps(func)
def wrapper(*args, **kwargs):
try:
func(*args, **kwargs)
except etcd3.exceptions.ConnectionFailedError as err:
raise etcd3.exceptions.ConnectionFailedError('etcd connection failed') from err
except etcd3.exceptions.ConnectionTimeoutError as err:
raise etcd3.exceptions.ConnectionTimeoutError('etcd connection timeout') from err
except Exception:
logger.exception('Some etcd error occurred')
return wrapper
class Etcd3Wrapper:
def __init__(self, *args, **kwargs):
self.client = etcd3.client(*args, **kwargs)
@readable_errors
def get(self, *args, value_in_json=False, **kwargs):
_value, _key = self.client.get(*args, **kwargs)
if _key is None or _value is None:
return None
return EtcdEntry(_key, _value, value_in_json=value_in_json)
@readable_errors
def put(self, *args, value_in_json=False, **kwargs):
_key, _value = args
if value_in_json:
_value = json.dumps(_value)
if not isinstance(_key, str):
_key = _key.decode('utf-8')
return self.client.put(_key, _value, **kwargs)
@readable_errors
def get_prefix(self, *args, value_in_json=False, **kwargs):
r = self.client.get_prefix(*args, **kwargs)
for entry in r:
e = EtcdEntry(*entry[::-1], value_in_json=value_in_json)
if e.value:
yield e
@readable_errors
def watch_prefix(self, key, timeout=0, value_in_json=False):
timeout_event = EtcdEntry(PseudoEtcdMeta(key=b'TIMEOUT'),
value=str.encode(json.dumps({'status': 'TIMEOUT',
'type': 'TIMEOUT'})),
value_in_json=value_in_json)
event_queue = queue.Queue()
def add_event_to_queue(event):
for e in event.events:
if e.value:
event_queue.put(EtcdEntry(e, e.value, value_in_json=value_in_json))
self.client.add_watch_prefix_callback(key, add_event_to_queue)
while True:
try:
while True:
v = event_queue.get(timeout=timeout)
yield v
except queue.Empty:
event_queue.put(copy.deepcopy(timeout_event))
class PsuedoEtcdEntry(EtcdEntry):
def __init__(self, key, value, value_in_json=False):
super().__init__(PseudoEtcdMeta(key=key.encode('utf-8')), value, value_in_json=value_in_json)