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
from time import time, sleep
from pprint import PrettyPrinter


class Dependecies:
    def __init__(self, deps):
        self.deps = deps

    @staticmethod
    def 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 find(self, service):
        all = set([service])
        if service in self.deps:
            services = self.deps[service]
            new = set(services) - all
            if new:
                print(f"{service}: {new}")
            all.update(services)            
            del self.deps[service]
            for one in new:
                all.update(self.find(one))
        return all


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


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 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) 
    return " ".join([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 dump_services(files):
    for filename in files:
        for service in dc_services(filename):
            print(service)


def dc_deps(dc_files, service):
    pp = PrettyPrinter(indent=2)
    for dc_file in dc_files:
        dep = Dependecies.load(dc_file)
        if service in dep:        
            all = sorted(Dependecies(dep).find(service))
            print(all)
            return
    print(f"Service {service} not found")



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


# AFM exec API: curl -fs -o /dev/null http://localhost:9001/actuator/health/readiness
# AQE: /opt/bin/aqe ping
# md-api-new: curl -fs -o /dev/null http://localhost:9008/actuator/health/readiness
# potgres: s6-setuidgid postgres; /usr/bin/pg_isready
# redis: /usr/bin/redis-cli ping
# result-cache: curl -fs -o /dev/null http://localhost:9041/actuator/health/readiness
# scan-model: curl -fs -o /dev/null http://localhost:9061/actuator/health/readiness
# sql-executor: curl -fs -o /dev/null http://localhost:9101/actuator/health/readiness


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


def check_pulsar():
    return curl_check("http://localhost:8080/admin/v2/tenants/public")


def check_postgres():
    return system("/usr/bin/pg_isready") == 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 up_and_wait(service, check, delay, duration):
    dc_run(["up", "-d", service])
    started = wait_for(check, delay, duration)
    echo(f"WAIT FOR {service}")
    if started:
        echo(f"READY {service}")
    else:
        echo(f"BAD {service}")


def start_tiger():
	up_and_wait("pulsar", check_pulsar, 5, 60)
	up(["redis", "jaeger", "router"])
	up_and_wait("postgres", check_postgres, 5, 60)
	up(["aqe", "afm-exec-api", "result-cache", "scan-model", "sql-executor", "metadata-api-new"])
	

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


DOCKER_COMPOSE = "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: up_and_wait("pulsar", check_pulsar, 10, 120),
    "postgres": lambda: up_and_wait("postgres", check_postgres, 10, 120),
    "data-loader": None, # co chci testovat? ze ma Exit code 0
}

DEFAULT_TIGER_SERVICES = [
    "afm-exec-api",
    "metadata-api-new",
    "aqe",
    "result-cache",
    "scan-model",
    "sql-executor"
]
TIGER_SERVICES = try_load_list("SERVICES", DEFAULT_TIGER_SERVICES)

COMMANDS = {
    "files": lambda args: dump_files(DC_FILES),
    "services": lambda args: dump_services(DC_FILES),
    "log": dc_log,
    "deps": lambda args: dc_deps(DC_FILES, args[0]),
    "wait": None,
    "up-wait": lambda args: SERVICES[args[0]](),
}


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