uncloud/etcd_wrapper.py

76 lines
2.4 KiB
Python
Raw Normal View History

2020-01-20 07:30:12 +00:00
import etcd3
import json
from functools import wraps
from uncloud import UncloudException
from uncloud.common import logger
class EtcdEntry:
2020-01-27 09:55:26 +00:00
def __init__(self, meta_or_key, value, value_in_json=True):
2020-01-20 07:30:12 +00:00
if hasattr(meta_or_key, 'key'):
# if meta has attr 'key' then get it
self.key = meta_or_key.key.decode('utf-8')
else:
# otherwise meta is the 'key'
self.key = meta_or_key
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:
return func(*args, **kwargs)
except etcd3.exceptions.ConnectionFailedError:
raise UncloudException('Cannot connect to etcd: is etcd running as configured in uncloud.conf?')
except etcd3.exceptions.ConnectionTimeoutError as err:
raise etcd3.exceptions.ConnectionTimeoutError('etcd connection timeout.') from err
2020-01-27 09:55:26 +00:00
except Exception as err:
logger.exception('Some etcd error occured. See syslog for details.', err)
2020-01-20 07:30:12 +00:00
return wrapper
class EtcdWrapper:
@readable_errors
def __init__(self, *args, **kwargs):
self.client = etcd3.client(*args, **kwargs)
@readable_errors
2020-01-27 09:55:26 +00:00
def get(self, *args, value_in_json=True, **kwargs):
2020-01-20 07:30:12 +00:00
_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
2020-01-27 09:55:26 +00:00
def put(self, *args, value_in_json=True, **kwargs):
2020-01-20 07:30:12 +00:00
_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
2020-01-27 09:55:26 +00:00
def get_prefix(self, *args, value_in_json=True, **kwargs):
2020-01-20 07:30:12 +00:00
event_iterator = self.client.get_prefix(*args, **kwargs)
for e in event_iterator:
yield EtcdEntry(*e[::-1], value_in_json=value_in_json)
@readable_errors
2020-01-27 09:55:26 +00:00
def watch_prefix(self, key, value_in_json=True):
2020-01-20 07:30:12 +00:00
event_iterator, cancel = self.client.watch_prefix(key)
for e in event_iterator:
if hasattr(e, '_event'):
2020-01-27 09:55:26 +00:00
e = getattr('e', '_event')
2020-01-20 07:30:12 +00:00
if e.type == e.PUT:
yield EtcdEntry(e.kv.key, e.kv.value, value_in_json=value_in_json)