119 lines
3.4 KiB
Python
119 lines
3.4 KiB
Python
import etcd3
|
|
import json
|
|
import queue
|
|
import copy
|
|
from uncloud import UncloudException
|
|
|
|
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:
|
|
return func(*args, **kwargs)
|
|
except etcd3.exceptions.ConnectionFailedError as err:
|
|
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
|
|
except Exception:
|
|
logger.exception(
|
|
"Some etcd error occured. See syslog for details."
|
|
)
|
|
|
|
return wrapper
|
|
|
|
|
|
class Etcd3Wrapper:
|
|
@readable_errors
|
|
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):
|
|
if hasattr(event, "events"):
|
|
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,
|
|
)
|