linux-workspace/bin/dcm

#!/usr/bin/env python3

"""
dcm down -v --remove-orphans
"""

import subprocess
import yaml
import pprint
import os
import re
from sys import argv, exit
from time import time, sleep
from pprint import PrettyPrinter
from operator import itemgetter


def echo(text):
    print(">> " + text)


def print_list(lst):
    for item in lst:
        print(item)


def error(text):
    print(text)
    exit(1)


def system(cmd):
    return subprocess.call(cmd, shell=True)


def split(text):
    return re.split(r'\s+', text)


def prefix_join(prefix, array):
    return " ".join([prefix + " " + item for item in array])


def join_args(args):
    return " ".join(args)


def dc_files(files):
    return prefix_join("-f", files)


def load_dc_file(dc_file):
    with open(dc_file) as yamlFile:
        dc = yaml.load(yamlFile, Loader=yaml.FullLoader)
    return dc


def dc_cmd(args):
    # split join
    args = join_args(args)
    return join_args([DOCKER_COMPOSE, DC_FILES_ARGS, args])


def dc_run(args):
    cmd = dc_cmd(args)
    print(cmd)
    return system(cmd)


def dc_log(args):
    args = ["logs", "--no-color"] + args + [r"| sed -r 's/^[^|]+\| //g'"]
    dc_run(args)


def dc_services(dc_file):
    dc = load_dc_file(dc_file)
    return [service for service in dc["services"].keys()]


def dc_commands(args):
    for key in sorted(COMMANDS.keys()):
        print(key)


def dump_services(files):
    services = [(filename, service) for filename in files for service in dc_services(filename)]
    services = map(itemgetter(1), services)
    #services = sorted(services)
    print_list(services)


def deps_load(yaml_filename):
    deps = {}
    with open(yaml_filename) as file:
        dc = yaml.load(file, Loader=yaml.FullLoader)
        for (name, conf) in dc['services'].items():
            depends_on = conf.get('depends_on', [])
            deps[name] = depends_on
    return deps


def deps_find(service, deps):
    all = set([service])
    if service in deps:
        services = deps[service]
        new = set(services) - all
        if new:
            new_list = " ".join(new)
            print(f"{service} >> {new_list}")
        all.update(services)
        del deps[service]
        for one in new:
            all.update(deps_find(one, deps))
    return all


def dc_deps(dc_files, args):
    service = args[0]
    pp = PrettyPrinter(indent=2)
    for dc_file in dc_files:
        deps = deps_load(dc_file)
        if service in deps:
            all = sorted(deps_find(service, deps))
            print("ALL >> " + " ".join(all))
            return
    print(f"Service {service} not found")


def dump_files(files):
    for file in files:
        print(file)


def curl_check(uri):
    cmd = "curl -f -s -o /dev/null " + uri
    echo(cmd)
    return system(cmd) == 0


def wait_for(check, delay, duration):
    start = time()
    while time() - start < duration:
        if check():
            return True
        else:
            echo(f"sleep for {delay} sec")
            sleep(delay)
    return False


def service_wait(service, check, delay = 10, duration = 120):
    echo(f"WAIT FOR {service}")
    started = wait_for(check, delay, duration)
    if started:
        echo(f"READY {service}")
    else:
        echo(f"BAD {service}")


def get_check(service):
    checker = SERVICES.get(service, None)
    if checker is None:
        error(f"No checker for {service}")
    else:
        return checker


def dc_check(args):
    for service in args:
        check = get_check(service)
        running = check()
        if running:
            print(f"{service} ON")
        else:
            print(f"{service} OFF")


def dc_up(services):
    dc_run(["up", "-d"] + services)


def dc_wait(args):
    for service in args:
        check = get_check(service)
        service_wait(service, check)


def dc_on(args):
    for service in args:
        if service == "tiger":
            tiger_on()
        else:
            check = get_check(service)
            dc_up([service])
            service_wait(service, check)


def dc_build_parallel(args):
    services = args
    dc_run(["build", "--parallel"] + args)


def tiger_on():
    dc_on(["pulsar", "postgres"])
    dc_up(["redis", "jaeger", "router"])
    dc_on(["result-cache", "metadata-api", "aqe", "sql-executor", "scan-model", "afm-exec-api"])


def try_load_list(name, default):
    if name in os.environ:
        return split(os.environ[name])
    else:
        return default


DOCKER_COMPOSE = os.environ.get("DC", "docker compose")

DEFAULT_DCF = ["docker-compose.yaml"]
DC_FILES = try_load_list("DCF", DEFAULT_DCF)
DC_FILES_ARGS = dc_files(DC_FILES)

SERVICES = {
    "pulsar":
        lambda:
            curl_check("http://localhost:8080/admin/v2/tenants/public"),
    "postgres":
        lambda:
            system("/usr/bin/pg_isready") == 0,
    "redis":
        lambda:
            system("/usr/bin/redis-cli ping") == 0,
    # AQE: /opt/bin/aqe ping
    "aqe":
        lambda:
            True,
    "metadata-api":
        lambda:
            curl_check("http://localhost:9008/actuator/health/readiness"),
    "afm-exec-api":
        lambda:
            curl_check("http://localhost:9001/actuator/health/readiness"),
    "result-cache":
        lambda:
            curl_check("http://localhost:9041/actuator/health/readiness"),
    "scan-model":
        lambda:
            curl_check("http://localhost:9061/actuator/health/readiness"),
    "sql-executor":
        lambda:
            curl_check("http://localhost:9101/actuator/health/readiness"),
    "data-loader":
        None, # co chci testovat? ze ma Exit code 0
}

COMMANDS = {
    "files": lambda args: dump_files(DC_FILES),
    "commands": dc_commands,
    "services": lambda args: dump_services(DC_FILES),
    "log": dc_log,
    "deps": lambda args: dc_deps(DC_FILES, args),
    "check": dc_check,
    "wait": dc_wait,
    "on": dc_on,
    "pb": dc_build_parallel,
    "bp": dc_build_parallel,
}


def main():
    args = argv[1:]
    action = args[0] if args else None
    if action in COMMANDS:
        COMMANDS[action](args[1:])
    else:
        exit(dc_run(args))


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        pass