etcd3_wrapper/etcd_wrapper.py

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)