1import os
  2import subprocess
  3import time
  4import tomllib
  5from importlib.util import find_spec
  6from pathlib import Path
  7
  8import click
  9
 10from plain.runtime import APP_PATH, settings
 11
 12from .poncho.manager import Manager as PonchoManager
 13from .utils import has_pyproject_toml
 14
 15
 16class ServicesPid:
 17    def __init__(self):
 18        self.pidfile = settings.PLAIN_TEMP_PATH / "dev" / "services.pid"
 19
 20    def write(self):
 21        pid = os.getpid()
 22        self.pidfile.parent.mkdir(parents=True, exist_ok=True)
 23        with self.pidfile.open("w+") as f:
 24            f.write(str(pid))
 25
 26    def rm(self):
 27        self.pidfile.unlink()
 28
 29    def exists(self):
 30        return self.pidfile.exists()
 31
 32
 33class Services:
 34    @staticmethod
 35    def get_services(root):
 36        if not has_pyproject_toml(root):
 37            return {}
 38
 39        with open(Path(root, "pyproject.toml"), "rb") as f:
 40            pyproject = tomllib.load(f)
 41
 42        return (
 43            pyproject.get("tool", {})
 44            .get("plain", {})
 45            .get("dev", {})
 46            .get("services", {})
 47        )
 48
 49    def __init__(self):
 50        self.poncho = PonchoManager()
 51
 52    @staticmethod
 53    def are_running():
 54        pid = ServicesPid()
 55        return pid.exists()
 56
 57    def run(self):
 58        # Each user of Services will have to check if it is running by:
 59        # - using the context manager (with Services())
 60        # - calling are_running() directly
 61        pid = ServicesPid()
 62        pid.write()
 63
 64        try:
 65            services = self.get_services(APP_PATH.parent)
 66            for name, data in services.items():
 67                env = {
 68                    **os.environ,
 69                    "PYTHONUNBUFFERED": "true",
 70                    **data.get("env", {}),
 71                }
 72                self.poncho.add_process(name, data["cmd"], env=env)
 73
 74            self.poncho.loop()
 75        finally:
 76            pid.rm()
 77
 78    def __enter__(self):
 79        if not self.get_services(APP_PATH.parent):
 80            # No-op if no services are defined
 81            return
 82
 83        if self.are_running():
 84            click.secho("Services already running", fg="yellow")
 85            return
 86
 87        print("Starting `plain dev services`")
 88        self._subprocess = subprocess.Popen(
 89            ["plain", "dev", "services"], cwd=APP_PATH.parent
 90        )
 91
 92        if find_spec("plain.models"):
 93            time.sleep(0.5)  # Give it a chance to hit on the first try
 94            subprocess.check_call(["plain", "models", "db-wait"], env=os.environ)
 95        else:
 96            # A bit of a hack to wait for the services to start
 97            time.sleep(3)
 98
 99    def __exit__(self, *args):
100        if not hasattr(self, "_subprocess"):
101            return
102
103        self._subprocess.terminate()
104
105        # Flush the buffer so the output doesn't spill over
106        self._subprocess.communicate()