diff options
| -rwxr-xr-x | Que/client.py | 120 |
1 files changed, 78 insertions, 42 deletions
diff --git a/Que/client.py b/Que/client.py index 6958576..f4ee601 100755 --- a/Que/client.py +++ b/Que/client.py @@ -5,7 +5,9 @@ simple client for que.run import argparse import configparser +import functools import http.client +import logging import os import subprocess import sys @@ -18,19 +20,62 @@ MAX_TIMEOUT = 99999999 # basically never timeout def auth(args): "Returns the auth key for the given ns from ~/.config/que.conf" + logging.debug("auth") namespace = args.target.split("/")[0] if namespace == "pub": return None conf_file = os.path.expanduser("~/.config/que.conf") if not os.path.exists(conf_file): sys.exit("you need a ~/.config/que.conf") - cfg = configparser.ConfigParser() - cfg.read(conf_file) + cfg = configparser.ConfigParser() + cfg.read(conf_file) return cfg[namespace]["key"] +def autodecode(bytestring): + """Attempt to decode bytes `bs` into common codecs, preferably utf-8. If + no decoding is available, just return the raw bytes. + + For all available codecs, see: + <https://docs.python.org/3/library/codecs.html#standard-encodings> + + """ + logging.debug("autodecode") + codecs = ["utf-8", "ascii"] + for codec in codecs: + try: + return bytestring.decode(codec) + except UnicodeDecodeError: + pass + return bytestring + + +def retry(exception, tries=4, delay=3, backoff=2): + "Decorator for retrying an action." + + def decorator(func): + @functools.wraps(func) + def func_retry(*args, **kwargs): + mtries, mdelay = tries, delay + while mtries > 1: + try: + return func(*args, **kwargs) + except exception as ex: + logging.debug(ex) + logging.debug("retrying...") + time.sleep(mdelay) + mtries -= 1 + mdelay *= backoff + return func(*args, **kwargs) + + return func_retry + + return decorator + + def send(args): "Send a message to the que." + logging.debug("send") key = auth(args) data = args.infile req = request.Request(f"{args.host}/{args.target}") @@ -38,6 +83,7 @@ def send(args): if key: req.add_header("Authorization", key) if args.serve: + logging.debug("serve") while not time.sleep(1): request.urlopen(req, data=data, timeout=MAX_TIMEOUT) @@ -45,19 +91,22 @@ def send(args): request.urlopen(req, data=data, timeout=MAX_TIMEOUT) -def recv(args): - "Receive a message from the que." +def then(args, msg): + "Perform an action when passed `--then`." + if args.then: + logging.debug("then") + subprocess.run( + args.then.replace(r"\msg", msg).replace(r"\que", args.target), + check=False, + shell=True, + ) - def _recv(_req): - msg = autodecode(_req.read()) - print(msg) - if args.then: - subprocess.run( - args.then.replace(r"\msg", msg).replace(r"que", args.target), - shell=True, - check=False, - ) +@retry(http.client.IncompleteRead, tries=10, delay=5, backoff=1) +@retry(http.client.RemoteDisconnected, tries=10, delay=2, backoff=2) +def recv(args): + "Receive a message from the que." + logging.debug("recv on: %s", args.target) params = urllib.parse.urlencode({"poll": args.poll}) req = request.Request(f"{args.host}/{args.target}?{params}") req.add_header("User-Agent", "Que/Client") @@ -66,32 +115,23 @@ def recv(args): req.add_header("Authorization", key) with request.urlopen(req) as _req: if args.poll: + logging.debug("poll") while not time.sleep(1): - _recv(_req) + logging.debug("reading") + msg = autodecode(_req.readline()) + logging.debug("read") + print(msg, end="") + then(args, msg) else: - _recv(_req) - - -def autodecode(bytestring): - """Attempt to decode bytes `bs` into common codecs, preferably utf-8. If - no decoding is available, just return the raw bytes. - - For all available codecs, see: - <https://docs.python.org/3/library/codecs.html#standard-encodings> - - """ - codecs = ["utf-8", "ascii"] - for codec in codecs: - try: - return bytestring.decode(codec) - except UnicodeDecodeError: - pass - return bytestring + msg = autodecode(_req.read()) + print(msg) + then(args, msg) def get_args(): "Command line parser" cli = argparse.ArgumentParser(description=__doc__) + cli.add_argument("--debug", action="store_true", help="log to stderr") cli.add_argument( "--host", default="http://que.run", help="where que-server is running" ) @@ -133,6 +173,12 @@ def get_args(): if __name__ == "__main__": ARGV = get_args() + if ARGV.debug: + logging.basicConfig( + format="%(asctime)s %(message)s", + level=logging.DEBUG, + datefmt="%Y.%m.%d..%H.%M.%S", + ) try: if ARGV.infile: send(ARGV) @@ -140,13 +186,3 @@ if __name__ == "__main__": recv(ARGV) except KeyboardInterrupt: sys.exit(0) - except urllib.error.HTTPError as err: - print(err) - sys.exit(1) - except http.client.RemoteDisconnected as err: - print("disconnected... retrying in 5 seconds") - time.sleep(5) - if ARGV.infile: - send(ARGV) - else: - recv(ARGV) |