spadinastan/tasks.py

import invoke
import jinja2
import yaml

import os
import pathlib
import typing

ns = invoke.Collection()

PYINFRA_COMMON_PREFIX = "pyinfra -vvv pyinfra/inventory.py"


def _get_tag(ctx: invoke.context.Context) -> str:
    """
    Gets the current tag, if it exists, for image versioning purposes.
    """
    try:
        out = ctx.run("git describe --tags --exact-match", hide=True)
        return out.stdout.strip()
    except:
        return "dev"


def _ensure_networks(ctx: invoke.context.Context):
    """
    Ensures that the required networks exist.
    """
    print("Ensuring networks exist...")
    output = ctx.run("docker network ls -f name=spad-internal -q", hide=True)
    if not bool(output.stdout):
        ctx.run("docker network create --scope=swarm spad-internal")
        print('\t✅ Created network "internal".')
    else:
        print('\t✅ Network "internal" already exists.')


def _ensure_swarm(ctx: invoke.context.Context):
    """
    Ensures that swarm mode is on.
    """
    print("Ensuring swarm mode is on...")
    try:
        ctx.run("docker swarm init --advertise-addr 192.168.1.17", hide=True)
    except:
        pass
    print("\t✅ Swarm in on.")


@invoke.task()
def system_updates(ctx):
    ctx.run(f"{PYINFRA_COMMON_PREFIX} pyinfra/system_updates.py")


@invoke.task()
def system_reboot(ctx):
    ctx.run(f"{PYINFRA_COMMON_PREFIX} pyinfra/reboot.py")


@invoke.task()
def generate_configuration(ctx):
    """
    Generates the service configuration file `services.yml`
    based on environment variables.

    To avoid polluting the environment, this can be called as
    `env $(cat .env | xargs) inv generate-configuration` where
    .env contains the variables.
    """
    template_loader = jinja2.FileSystemLoader(searchpath="./")
    template_env = jinja2.Environment(
        loader=template_loader, undefined=jinja2.StrictUndefined
    )
    template = template_env.get_template("services.yml.j2")

    with open("services.yml", "w") as outfile:
        outfile.write(template.render(**os.environ))


@invoke.task()
def start(ctx, services: str):
    """
    Starts one of more services, defined by a comma-separated list of labels.

    Services should correspond to an entry in `services.yml`.
    """
    _ensure_swarm(ctx)
    _ensure_networks(ctx)

    current_version = _get_tag(ctx)

    print("Starting services...")

    for service in services.split(","):
        with open("services.yml", "r") as config:
            service_config = yaml.load(config, Loader=yaml.Loader)["services"][service]

        ctx.run(
            f"docker build -t spadinaistan-{service}:{current_version} -f services/Dockerfile-{service} .",
            hide=True,
        )
        ports_args = (
            f"--publish {ports}\\\n" for ports in service_config.get("ports", [])
        )

        volume_args = []

        for volume in service_config.get("volumes", []):
            source, target = volume.split(":")
            volume_args.append(f"--mount type=bind,source={source},target={target}\\\n")

        env_files = (
            f"--env-file {env_file}\\\n"
            for env_file in service_config.get("env_files", [])
        )

        try:
            out = ctx.run(
                f"""docker service create \
                            --name spad-{service}\
                            --network spad-internal\
                            {" ".join(env_files)}\
                            --hostname=\"{service}\"\
                            {" ".join(volume_args)}\
                            {" ".join(ports_args)}\
                            spadinaistan-{service}:{current_version}""",
                hide=True,
            )
        except:
            print(out.stdout)
            print(out.stderr)
            continue

        print(f"\t✅ {service} is running.")


@invoke.task()
def stop(ctx: invoke.context.Context, services: str):
    """
    Stops the provided list of services, as a comma-separated list
    of labels.
    """
    print("Stopping services...")
    for service in services.split(","):
        ctx.run(f"docker service rm spad-{service}", hide=True)
        print(f"\t✅ {service} is stopped.")


services = invoke.Collection("service")
services.add_task(start)
services.add_task(stop)

server = invoke.Collection("server")

server.add_task(system_updates, name="update")
server.add_task(system_reboot, name="reboot")

ns.add_collection(server)
ns.add_collection(services)
ns.add_task(generate_configuration)