#!/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): # noqa: D213 """Get the health of containers of a Compose 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): # noqa: D213 """Returns a list (could be empty) of containers that are not healthy. Or in the case that a health check is not set, containers that are not running. """ 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): # noqa: D213 """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) if not healths: raise RuntimeError(f"No containers found for project {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 ) try: unhealthy = wait_for_project_health(project, args.timeout) except Exception as e: # pylint: disable=broad-except arg_parser.error(str(e)) 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())