compose-health-check

#!/usr/bin/env python3
# pylint: disable=invalid-name
"""Checks the health of a Docker Compose deployment."""


import argparse
import datetime
import os
import sys
import time
import docker  # pylint: disable=import-error
import dotenv  # pylint: disable=import-error


def get_project_containers_health(client, project):
    """Get a Docker connection and the Docker Compose project name and return a
    dictionary of container objects and their health status (if available, else
    their status)."""
    return {
        i: i.attrs["State"]["Health"]["Status"]
        if "Health" in i.attrs["State"]
        else i.attrs["State"]["Status"]
        for i in client.containers.list(
            filters={"label": f"com.docker.compose.project={project}"}
        )
    }


def get_unhealthy_project_containers(client, project):
    """Returns a list (could be empty) of containers that are not healthy or
    not running (in case a health check is not set)."""
    healths = get_project_containers_health(client, project)
    return [x for x in healths if healths[x] not in ["healthy", "running"]]


def wait_for_project_health(project, timeout):
    """Wait for a Docker Compose project to be healthy.

    If all of the containers are healthy or running (because there's no
    healthcheck set), return an empty list. If any of the containers are
    starting (health check hasn't run yet or in the grace period), sleep 10 and
    try again (ignoring the timeout). Otherwise, if the timeout hasn't passed,
    sleep 10 and try again, if over the timeout, return a list of the
    containers that aren't healthy or aren't running.
    """
    starttime = datetime.datetime.now()
    deadline = starttime + datetime.timedelta(seconds=timeout)
    docker_client = docker.from_env()
    docker_client.ping()
    while True:
        healths = get_project_containers_health(docker_client, project)
        # pylint: disable=no-else-return
        if not get_unhealthy_project_containers(docker_client, project):
            return []
        elif any(map(lambda x: x == "starting", healths.values())):
            time.sleep(10)
            continue
        elif datetime.datetime.now() > deadline:
            break
        else:
            time.sleep(10)

    return get_project_containers_health(docker_client, project)


def get_project_name(args):
    """Returns the name of the Docker Compose project, or None if unable to."""
    if args.project:
        return args.project
    env = dotenv.dotenv_values(dotenv_path=".env")
    if "COMPOSE_PROJECT_NAME" in env:
        return env["COMPOSE_PROJECT_NAME"]
    if "COMPOSE_PROJECT_NAME" in os.environ:
        return os.environ["COMPOSE_PROJECT_NAME"]
    return None


def main():
    """Main entrypoint."""

    epilog = (
        "The Docker Compose project name is resolved in the following order:"
        "\n1. From the command line parameter."
        "\n2. From the COMPOSE_PROJECT_NAME variable in the .env file."
        "\n3. From the COMPOSE_PROJECT_NAME environment variable."
    )
    arg_parser = argparse.ArgumentParser(
        description=__doc__,
        epilog=epilog,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    arg_parser.add_argument(
        "project", help="Name of the Compose Project.", nargs="?"
    )
    arg_parser.add_argument(
        "-t",
        "--timeout",
        help="Number of seconds to wait for the deployment to be healthy, defaults to 300.",  # noqa: E501
        type=int,
        default=300,
    )
    args = arg_parser.parse_args()

    project = get_project_name(args)
    if project is None:
        arg_parser.error(
            "Compose project wasn't specified, the COMPOSE_PROJECT_NAME variable is missing from the environment and from the .env file."  # noqa: E501
        )

    unhealthy = wait_for_project_health(project, args.timeout)

    if not unhealthy:
        print(f"Project {project} is healthy.")
        return 0
    print(
        f"""The {", ".join(map(lambda x: x.name, unhealthy))} container for project {project} are not healthy."""  # noqa: E501
    )
    return 1


if __name__ == "__main__":
    sys.exit(main())