goat_monitor/goat_monitor.py

# %% imports
import subprocess
from pathlib import Path
from typing import List

import click
import gotify
import numpy as np
import toml

__version__ = "v0.0.0"


# %% commands
@click.command()
@click.argument("command", nargs=-1, required=False, type=str)
@click.option(
    "--config",
    default="~/.config/goat_monitor.toml",
    type=click.Path(path_type=Path),
    help="Use a .toml configuration file",
)
@click.option(
    "--retries",
    default=0,
    type=int,
    help="Number of times to try re-running failed command. -1 to retry indefinitely until success",
)
def wrap(command: List[str], config: Path, retries: int):
    """Wrap an arbitrary command with gotify notifications"""

    # read settings first
    with open(config, "r") as f:
        settings = toml.load(f)

    # gotify configuration
    url = settings["server"]
    app_token = settings["app_token"]
    with gotify.Gotify(base_url=url, app_token=app_token) as gotify_connection:
        # this is just to test configuration.
        # I don't know how long a gotify connection will stay up so rather than thinking about it
        # I'm just creating a new one whenever I need to send a message
        pass

    if retries == -1:
        retries = np.inf
    if retries < 0:
        raise ValueError("Invalid number of retries specified")

    for attempt in range(retries + 1):
        # run the command
        result = subprocess.run(
            " ".join(command),
            shell=True,
            stderr=subprocess.STDOUT,
            stdout=subprocess.PIPE,
            encoding="utf-8",
        )

        # TODO: print lines real time as the subprocess runs
        print(result.stdout, end="")

        if result.returncode:
            # failed
            title = f"Command failed with exit code {result.returncode}"
            if attempt < retries:
                title += f" - Retrying (attempt {attempt+1}/{retries})"
            else:
                title += " - Aborting"

        else:
            title = "Command succeeded"

        MAX_LINES = 20
        lines = result.stdout.splitlines()
        if len(lines) > MAX_LINES:
            lines = lines[-MAX_LINES:]
        message = (
            "Command:\n"
            + " ".join(command)
            + f'\n\nResult{ " (truncated)" if len(lines) > MAX_LINES else ""}:\n'
            + "\n".join(lines)
        )

        with gotify.Gotify(base_url=url, app_token=app_token) as gotify_connection:
            gotify_connection.create_message(message=message, title=title)

    return result.returncode


# %% main
if __name__ == "__main__":
    wrap()