find ucloud -name \*.py -exec sed -i "s/ucloud/uncloud/g" {} \;

uncloud/common/etcd_wrapper.py

@@ -0,0 +1,117 @@
+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:
+            return 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 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,
+        )