102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
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)
|