import etcd3 import json import signal import threading import time import queue from collections import namedtuple from dataclasses import dataclass from datetime import datetime from os import path PseudoEtcdMeta = namedtuple("PseudoEtcdMeta", ["key"]) @dataclass(init=False) class EtcdEntry: 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) key: str value: str class Etcd3Wrapper(object): def __init__(self, *args, **kwargs): self.client = etcd3.client(*args, **kwargs) def get(self, *args, value_in_json=False, **kwargs): _value, _key = self.client.get(*args, **kwargs) if _key is None: return None return EtcdEntry(_key, _value, value_in_json=value_in_json) 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) def get_prefix(self, *args, value_in_json=False, **kwargs): r = self.client.get_prefix(*args, **kwargs) for entry in r: yield EtcdEntry(*entry[::-1], value_in_json=value_in_json) # def watch_prefix(self, *args, timeout=0, value_in_json=False, **kwargs): # def raise_timeout(signum, frame): # raise TimeoutError("timeout") # r, _ = self.client.watch_prefix(*args, **kwargs) # if timeout > 0: # signal.signal(signal.SIGALRM, raise_timeout) # while True: # try: # signal.alarm(timeout) # for event in r: # signal.alarm(0) # if event.value: # event = EtcdEntry(event, event.value, value_in_json=value_in_json) # signal.alarm(timeout) # yield event # else: # signal.alarm(timeout) # except TimeoutError: # _value = {"status": "TIMEOUT"} # yield EtcdEntry(PseudoEtcdMeta(key=b"TIMEOUT"), value=str.encode(json.dumps(_value))) def watch_prefix(self, key_prefix, timeout, value_in_json=False, **kwargs): q = queue.Queue() r, _ = self.client.watch_prefix(key_prefix, **kwargs) t = threading.Thread(target=watch_prefix_generator, args=(value_in_json, r, q)) t.start() while True: try: v = q.get(timeout=timeout) yield v except queue.Empty: _value = {"status": "TIMEOUT"} e = EtcdEntry(PseudoEtcdMeta(key=b"TIMEOUT"), value=str.encode(json.dumps(_value))) yield e def watch_prefix_generator(value_in_json, generator, q): for event in generator: if event.value: event = EtcdEntry(event, event.value) q.put(event)