Bugfix for helper CLI, expose client structs, implement helpers on process state

Author: timkpaine

.gitignore

@@ -143,5 +143,8 @@ _template/labextension
 
 # Supervisor
 supervisord.pid
+supervisor-2024*.cfg
+pydantic.json
+Untitled*.ipynb
 tmp
 

Makefile

@@ -54,7 +54,9 @@ annotate:  ## run python type annotation checks with mypy
 
 test:  ## run python tests
 	python -m pytest -v airflow_supervisor/tests --junitxml=junit.xml
-	pgrep -iaf supervisord | xargs kill -15
+
+kill:
+	bash -c "pgrep -iaf supervisord | xargs kill -15"
 
 coverage:  ## run tests and collect test coverage
 	python -m pytest -v airflow_supervisor/tests --junitxml=junit.xml --cov=airflow_supervisor --cov-branch --cov-fail-under=60 --cov-report term-missing --cov-report xml

airflow_supervisor/airflow/commands.py

@@ -39,12 +39,18 @@ def _wait_or_while(until, timeout: int = 5) -> bool:
     return False
 
 
+def _raise_or_exit(val: bool, exit: bool):
+    if exit:
+        raise Exit(int(not val))
+    return val
+
+
 def write_supervisor_config(cfg_json: str, _exit: Annotated[bool, Argument(hidden=True)] = True):
     cfg_obj = SupervisorAirflowConfiguration.model_validate_json(cfg_json)
     if not _check_same(cfg_obj):
         log.critical("Configs don't match")
     cfg_obj._write_self()
-    return Exit(0) if _exit else True
+    return _raise_or_exit(True, _exit)
 
 
 def start_supervisor(
@@ -57,13 +63,13 @@ def start_supervisor(
     if not _check_same(cfg_obj):
         log.critical("Configs don't match")
     if _check_running(cfg_obj):
-        return Exit(0) if _exit else True
+        return _raise_or_exit(True, _exit)
     cfg_obj.start(daemon=True)
     running = _wait_or_while(until=lambda: cfg_obj.running(), timeout=30)
     if not running:
         log.critical("Still not running 30s after start command!")
-        return Exit(1) if _exit else False
-    return Exit(0) if _exit else True
+        return _raise_or_exit(False, _exit)
+    return _raise_or_exit(True, _exit)
 
 
 def start_programs(
@@ -77,21 +83,31 @@ def start_programs(
     # TODO
     ret = client.startAllProcesses()
     log.info(ret)
-    return Exit(0) if _exit else True
+    return _raise_or_exit(True, _exit)
 
 
 def check_programs(
     cfg: Annotated[
         Path, Option(exists=True, file_okay=True, dir_okay=False, writable=False, readable=True, resolve_path=True)
     ],
+    check_running: bool = False,
     _exit: Annotated[bool, Argument(hidden=True)] = True,
 ):
     cfg_obj = SupervisorAirflowConfiguration.model_validate_json(cfg.read_text())
     client = SupervisorRemoteXMLRPCClient(cfg=cfg_obj)
     # TODO
     ret = client.getAllProcessInfo()
-    log.info(ret)
-    return Exit(0) if _exit else True
+    for r in ret:
+        log.info(r.model_dump_json())
+    if check_running:
+        meth = "running"
+    else:
+        meth = "ok"
+    if all(getattr(p, meth)() for p in ret):
+        log.info("all processes ok")
+        return _raise_or_exit(True, _exit)
+    log.info("processes not ok")
+    return _raise_or_exit(False, _exit)
 
 
 def restart_programs(
@@ -107,7 +123,7 @@ def restart_programs(
     log.info(ret1)
     ret2 = client.startAllProcesses()
     log.info(ret2)
-    return Exit(0) if _exit else True
+    return _raise_or_exit(True, _exit)
 
 
 def stop_supervisor(
@@ -121,8 +137,8 @@ def stop_supervisor(
     not_running = _wait_or_while(until=lambda: not cfg_obj.running(), timeout=30)
     if not not_running:
         log.critical("Still running 30s after stop command!")
-        return Exit(1) if _exit else False
-    return Exit(0) if _exit else True
+        return _raise_or_exit(False, _exit)
+    return _raise_or_exit(True, _exit)
 
 
 def kill_supervisor(
@@ -136,8 +152,8 @@ def kill_supervisor(
     still_running = _wait_or_while(until=lambda: not cfg_obj.running(), timeout=30)
     if still_running:
         log.critical("Still running 30s after kill command!")
-        return Exit(1) if _exit else False
-    return Exit(0) if _exit else True
+        return _raise_or_exit(False, _exit)
+    return _raise_or_exit(True, _exit)
 
 
 def remove_supervisor_config(
@@ -152,14 +168,14 @@ def remove_supervisor_config(
         still_running = kill_supervisor(cfg_obj, _exit=False)
 
     if still_running:
-        return Exit(1) if _exit else True
+        return _raise_or_exit(False, _exit)
 
     # TODO move to config
     sleep(5)
 
     # TODO make optional
     cfg_obj.rmdir()
-    return Exit(0) if _exit else True
+    return _raise_or_exit(True, _exit)
 
 
 def _add_to_typer(app, command: _SupervisorTaskStep, foo):

airflow_supervisor/client/__init__.py

@@ -1 +1 @@
-from .xmlrpc import SupervisorRemoteXMLRPCClient
+from .xmlrpc import *

airflow_supervisor/client/xmlrpc.py

@@ -6,10 +6,10 @@
 
 from ..config import SupervisorAirflowConfiguration
 
-__all__ = ("SupervisorRemoteXMLRPCClient",)
+__all__ = ("SupervisorRemoteXMLRPCClient", "ProcessState", "SupervisorState", "SupervisorMethodResult", "ProcessInfo")
 
 
-class State(Enum):
+class ProcessState(Enum):
     STOPPED = 0
     STARTING = 10
     RUNNING = 20
@@ -66,7 +66,7 @@ class SupervisorMethodResult(Enum):
 class ProcessInfo(BaseModel):
     name: str
     group: str
-    state: State
+    state: ProcessState
     description: str
     start: datetime
     stop: datetime
@@ -78,6 +78,17 @@ class ProcessInfo(BaseModel):
     stderr_logfile: str
     pid: int
 
+    def running(self):
+        return self.state in (ProcessState.STARTING, ProcessState.RUNNING, ProcessState.STOPPING)
+
+    def ok(self):
+        return self.state in (
+            ProcessState.STARTING,
+            ProcessState.RUNNING,
+            ProcessState.STOPPING,
+            ProcessState.STOPPED,
+        ) or (self.state == ProcessState.EXITED and self.exitstatus == 0)
+
 
 class SupervisorRemoteXMLRPCClient(object):
     """A light wrapper over the supervisor xmlrpc api: http://supervisord.org/api.html"""

airflow_supervisor/config/supervisor.py

@@ -87,20 +87,18 @@ def _setup_config_and_working_dir(self):
             else:
                 self.working_dir = tempdir / f"supervisor-{now}"
                 self.supervisord.directory = self.working_dir
-
-            # force pidfile to be in working dir if not otherwise set
-            if not self.supervisord.pidfile:
-                self.supervisord.pidfile = self.working_dir / "supervisord.pid"
-
-            # force logfile to be in working dir if not otherwise set
-            if not self.supervisord.logfile:
-                self.supervisord.logfile = self.working_dir / "supervisord.log"
-
             using_default_working_dir = True
-
         else:
             using_default_working_dir = False
 
+        # force pidfile to be in working dir if not otherwise set
+        if not self.supervisord.pidfile:
+            self.supervisord.pidfile = self.working_dir / "supervisord.pid"
+
+        # force logfile to be in working dir if not otherwise set
+        if not self.supervisord.logfile:
+            self.supervisord.logfile = self.working_dir / "supervisord.log"
+
         if self.config_path == "":
             if using_default_working_dir:
                 self.config_path = self.working_dir / "supervisor.cfg"

airflow_supervisor/tests/test_client.py airflow_supervisor/tests/client/test_client_e2e.py

@@ -0,0 +1,30 @@
+from datetime import datetime
+
+from airflow_supervisor.client import ProcessInfo, ProcessState
+
+
+def _gen() -> ProcessInfo:
+    return ProcessInfo(
+        name="test",
+        group="test",
+        state=ProcessState.UNKNOWN,
+        description="",
+        start=datetime.now(),
+        stop=datetime.now(),
+        now=datetime.now(),
+        spawner="",
+        exitstatus=0,
+        logfile="",
+        stdout_logfile="",
+        stderr_logfile="",
+        pid=0,
+    )
+
+
+def test_ok():
+    x = _gen()
+    x.state = ProcessState.RUNNING
+    assert x.ok()
+    x.state = ProcessState.EXITED
+    x.exitstatus = 0
+    assert x.ok()

airflow_supervisor/tests/client/test_client_structs.py

@@ -0,0 +1,18 @@
+from pathlib import Path
+
+from airflow_supervisor import AirflowConfiguration, ProgramConfiguration, SupervisorAirflowConfiguration
+
+if __name__ == "__main__":
+    path = Path(__file__).parent.parent.parent
+    cfg = SupervisorAirflowConfiguration(
+        airflow=AirflowConfiguration(port="*:9090"),
+        working_dir=path,
+        path=path,
+        program={
+            "test": ProgramConfiguration(
+                command="sleep 1 && exit 1",
+            )
+        },
+    )
+    print(cfg._pydantic_path)
+    cfg._write_self()