You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

102 lines
3.1 KiB

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"])
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))
while True:
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)