upload.py

#!/usr/bin/env python
import os
import json
import cherrypy
import time
import threading
from cklib.graph import Graph
from cklib.graph.export import node_from_dict
from prometheus_client import Summary
from prometheus_client.exposition import generate_latest, CONTENT_TYPE_LATEST
from cklib.args import ArgumentParser
from cklib.logging import log
from signal import signal, SIGTERM, SIGINT


shutdown_event = threading.Event()

metrics_upload_graph = Summary(
    "upload_graph_seconds",
    "Time it took the upload() function",
)


def handler(sig, frame) -> None:
    log.info("Shutting down")
    shutdown_event.set()


def main() -> None:
    signal(SIGINT, handler)
    signal(SIGTERM, handler)
    arg_parser = ArgumentParser(description="Cloudkeeper Webserver Test")
    WebServer.add_args(arg_parser)
    arg_parser.parse_args()
    web_server = WebServer(WebApp())
    web_server.daemon = True
    web_server.start()
    shutdown_event.wait()
    web_server.shutdown()
    os._exit(0)


class WebApp:
    def __init__(self) -> None:
        self.mountpoint = ArgumentParser.args.web_path
        local_path = os.path.abspath(os.path.dirname(__file__))
        config = {
            "tools.gzip.on": True,
            #            "tools.staticdir.index": "index.html",
            #            "tools.staticdir.on": True,
            #            "tools.staticdir.dir": f"{local_path}/static",
        }
        self.config = {"/": config}
        if self.mountpoint not in ("/", ""):
            self.config[self.mountpoint] = config

    @cherrypy.expose
    @cherrypy.tools.allow(methods=["GET"])
    def health(self):
        cherrypy.response.headers["Content-Type"] = "text/plain"
        return "ok\r\n"

    @cherrypy.expose
    @cherrypy.tools.allow(methods=["GET"])
    def metrics(self):
        cherrypy.response.headers["Content-Type"] = CONTENT_TYPE_LATEST
        return generate_latest()

    @cherrypy.config(**{"response.timeout": 3600})
    @cherrypy.expose()
    @cherrypy.tools.allow(methods=["POST"])
    @metrics_upload_graph.time()
    def upload(self):
        log.info("Receiving Graph data")
        start_time = time.time()
        g = Graph()
        node_map = {}
        for line in cherrypy.request.body.readlines():
            line = line.decode("utf-8")
            try:
                data = json.loads(line)
                if "id" in data:
                    node = node_from_dict(data)
                    node_map[data["id"]] = node
                    g.add_node(node)
                elif (
                    "from" in data
                    and "to" in data
                    and data["from"] in node_map
                    and data["to"] in node_map
                ):
                    g.add_edge(node_map[data["from"]], node_map[data["to"]])
            except json.decoder.JSONDecodeError:
                continue
        log.info(
            f"Received Graph with {g.number_of_nodes()} nodes"
            f" and {g.number_of_edges()} edges"
            f" in {time.time() - start_time:.2f} seconds."
        )
        cherrypy.response.headers["Content-Type"] = "text/plain"
        return f"ok\r\n"


class WebServer(threading.Thread):
    def __init__(self, webapp) -> None:
        super().__init__()
        self.name = "webserver"
        self.webapp = webapp

    @property
    def serving(self):
        return cherrypy.engine.state == cherrypy.engine.states.STARTED

    def run(self) -> None:
        # CherryPy always prefixes its log messages with a timestamp.
        # The next line monkey patches that time method to return a
        # fixed string. So instead of having duplicate timestamps in
        # each web server related log message they are now prefixed
        # with the string 'CherryPy'.
        cherrypy._cplogging.LogManager.time = lambda self: "CherryPy"
        cherrypy.engine.unsubscribe("graceful", cherrypy.log.reopen_files)

        # We always mount at / as well as any user configured --web-path
        cherrypy.tree.mount(
            self.webapp,
            "",
            self.webapp.config,
        )
        if self.webapp.mountpoint not in ("/", ""):
            cherrypy.tree.mount(
                self.webapp,
                self.webapp.mountpoint,
                self.webapp.config,
            )
        cherrypy.config.update(
            {
                "global": {
                    "engine.autoreload.on": False,
                    "server.socket_host": ArgumentParser.args.web_host,
                    "server.socket_port": ArgumentParser.args.web_port,
                    "server.max_request_body_size": 1048576000,  # 1GB
                    "server.socket_timeout": 60,
                    "log.screen": False,
                    "log.access_file": "",
                    "log.error_file": "",
                    "tools.log_headers.on": False,
                    "tools.encode.on": True,
                    "tools.encode.encoding": "utf-8",
                    "request.show_tracebacks": False,
                    "request.show_mismatched_params": False,
                }
            }
        )
        cherrypy.engine.start()
        cherrypy.engine.block()

    def shutdown(self):
        log.debug("Received request to shutdown http server threads")
        cherrypy.engine.exit()

    @staticmethod
    def add_args(arg_parser: ArgumentParser) -> None:
        arg_parser.add_argument(
            "--web-port",
            help="Web Port (default 8000)",
            default=8000,
            dest="web_port",
            type=int,
        )
        arg_parser.add_argument(
            "--web-host",
            help="IP to bind to (default: ::)",
            default="::",
            dest="web_host",
            type=str,
        )
        arg_parser.add_argument(
            "--web-path",
            help="Web root in browser (default: /)",
            default="/",
            dest="web_path",
            type=str,
        )


if __name__ == "__main__":
    main()