Slight tweak to use API keys for health check on recent versions of GitLab

[kfiresmith]

After upgrading to GitLab 16, our checks broke. Turns out we need to use API keys now as part of the health check endpoint URL.

I got a minimally working version of the check set up and tested. It needs further work from an actual Python coder to ensure error checking, input validation, and testing for corner cases if any.

Because while it works but it's not a full feature add by my standards, I'm not going to do a PR, but I'll attach the new check here for anyone else's benefit who comes looking.

[kfiresmith]

#!/usr/bin/env python3

import json
import socket
import sys
import urllib.request

program_name = "check_gitlab"

def get_readiness(server, token):
    if (not server.endswith("/", 0)):
        server = server + "/"

    request = urllib.request.urlopen(server + "-/readiness?all=1&token=" + token)
    responseStr = str(request.read().decode("utf-8"))

    return json.loads(responseStr)

def check_module(readiness_json, module):
    try:
        status = readiness_json[module][0]["status"]
    except:
        status = "unknown"

    try:
        message = readiness_json[module][0]["message"]
    except:
        message = ""

    return status, message

def check_all(readiness_json):
    results = dict()
    results["cache_check"] = check_module(readiness_json, "cache_check")
    results["db_check"] = check_module(readiness_json, "db_check")
    results["gitaly_check"] = check_module(readiness_json, "gitaly_check")
    results["queues_check"] = check_module(readiness_json, "queues_check")
    results["rate_limiting_check"] = check_module(readiness_json, "rate_limiting_check")
    results["redis_check"] = check_module(readiness_json, "redis_check")
    results["repository_cache_check"] = check_module(readiness_json, "repository_cache_check")
    results["sessions_check"] = check_module(readiness_json, "sessions_check")
    results["shared_state_check"] = check_module(readiness_json, "shared_state_check")
    results["trace_chunks_check"] = check_module(readiness_json, "trace_chunks_check")

    return results

def parse_arguments(argv):
    argument_map = dict()
    argument_map["-h"] = 0
    argument_map["-s"] = ""
    argument_map["--token"] = ""
    argument_map["--check-all"] = 0
    argument_map["--check-cache"] = 0
    argument_map["--check-db"] = 0
    argument_map["--check-gitaly"] = 0
    argument_map["--check-queues"] = 0
    argument_map["--check-rate-limiting"] = 0
    argument_map["--check-redis"] = 0
    argument_map["--check-repository-cache"] = 0
    argument_map["--check-sessions"] = 0
    argument_map["--check-shared-state"] = 0
    argument_map["--check-trace-chunks"] = 0

    prev_argument = ""

    for argument in argv[1::]:
        if (argument.startswith("-")):
            argument_map[argument] = 1
        elif (prev_argument.startswith("-")):
            argument_map[prev_argument] = argument
        else:
            print_usage()
            sys.exit(3)

        prev_argument = argument

    if (len(argument_map) != 14):
        print_usage()
        sys.exit(3)

    #print(argument_map)
    return argument_map

def print_usage():
    print("Usage: " + program_name + " -s <server_url> --token <API token> [-h] [--check-all] [--check-cache] [--check-db] [--check-gitaly] [--check-queues] [--check-rate-limiting] [--check-redis] [--check-repository-cache] [--check-sessions] [--check-shared-state] [--check-trace-chunks]")

def print_help():
    print("Copyright (c) 2019 Kevin Krumbiegel")
    print("\n")
    print_usage();
    print("\n");
    print("-s <server_url>       = URL of the server to be checked. Must include trailing slash. (i.e. https://gitlab.example.org/)")
    print("--token <API token>   = Unique API token for Gitlab health checks (i.e. 34pfFi6T8dCj1G8zh2cs")
    print("-h                    = This screen")
    print("Checks:")
    print("--check-all           = Enable all checks")
    print("--check-cache         = Enable checking cache status")
    print("--check-db            = Enable checking database status")
    print("--check-gitaly        = Enable checking gitaly status")
    print("--check-queues        = Enable checking queues status")
    print("--check-rate-limiting = Enable checking rate limiting status")
    print("--check-redis         = Enable checking redis status")
    print("--check-repository-cache = Enable checking repository cache status")
    print("--check-sessions        = Enable checking sessions status")
    print("--check-shared-state  = Enable checking shared state status")
    print("--check-trace-chunks        = Enable checking trace chunks status")
    print("\n")

def run_checks(arguments):
    checks = check_all(get_readiness(arguments["-s"], arguments["--token"]))
    result = 0
    checks_done = 0

    if arguments["--check-cache"] or arguments["--check-all"]:
        checks_done += 1
        if str(checks["cache_check"][0]) != "ok":
            print("Cache check: Status=" + str(checks["cache_check"][0]) + ", Message=" + str(checks["cache_check"][1]))
            result = 2

    if arguments["--check-db"] or arguments["--check-all"]:
        checks_done += 1
        if str(checks["db_check"][0]) != "ok":
            print("Database check: Status=" + str(checks["db_check"][0]) + ", Message=" + str(checks["db_check"][1]))
            result = 2

    if arguments["--check-gitaly"] or arguments["--check-all"]:
        checks_done += 1
        if str(checks["gitaly_check"][0]) != "ok":
            print("Gitaly check: Status=" + str(checks["gitaly_check"][0]) + ", Message=" + str(checks["gitaly_check"][1]))
            result = 2

    if arguments["--check-queues"] or arguments["--check-all"]:
        checks_done += 1
        if str(checks["queues_check"][0]) != "ok":
            print("Queues check: Status=" + str(checks["queues_check"][0]) + ", Message=" + str(checks["queues_check"][1]))
            result = 2

    if arguments["--check-rate-limiting"] or arguments["--check-all"]:
        checks_done += 1
        if str(checks["rate_limiting_check"][0]) != "ok":
            print("Rate-limiting check: Status=" + str(checks["rate_limiting_check"][0]) + ", Message=" + str(checks["rate_limiting_check"][1]))
            result = 2
            
    if arguments["--check-redis"] or arguments["--check-all"]:
        checks_done += 1
        if str(checks["redis_check"][0]) != "ok":
            print("Redis check: Status=" + str(checks["redis_check"][0]) + ", Message=" + str(checks["redis_check"][1]))
            result = 2

    if arguments["--check-repository-cache"] or arguments["--check-all"]:
        checks_done += 1
        if str(checks["repository_cache_check"][0]) != "ok":
            print("Repository cache check: Status=" + str(checks["repository_cache_check"][0]) + ", Message=" + str(checks["repository_cache_check"][1]))
            result = 2

    if arguments["--check-sessions"] or arguments["--check-all"]:
        checks_done += 1
        if str(checks["sessions_check"][0]) != "ok":
            print("Sessions check: Status=" + str(checks["sessions_check"][0]) + ", Message=" + str(checks["sessions_check"][1]))
            result = 2
            
    if arguments["--check-trace-chunks"] or arguments["--check-all"]:
        checks_done += 1
        if str(checks["trace_chunks_check"][0]) != "ok":
            print("Trace chunks check: Status=" + str(checks["trace_chunks_check"][0]) + ", Message=" + str(checks["trace_chunks_check"][1]))
            result = 2
            
    if arguments["--check-shared-state"] or arguments["--check-all"]:
        checks_done += 1
        if str(checks["shared_state_check"][0]) != "ok":
            print("Shared state check: Status=" + str(checks["shared_state_check"][0]) + ", Message=" + str(checks["shared_state_check"][1]))
            result = 2

    if (result == 0):
        print("OK - All checks ok")

    print(str(checks_done) + " checks done.")
    
    return result


def run():
    arguments = parse_arguments(sys.argv)
    if (arguments["-h"] != 0):
        print_help();
        sys.exit(0)

    if (arguments["-s"] == ""):
        print_usage();
        sys.exit(3)

    sys.exit(run_checks(arguments))

try:
    run()
except Exception as e:
    print(e)
    sys.exit(3)