98 lines
3.0 KiB
Python
98 lines
3.0 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)
|