@@ -39,12 +39,18 @@ def _wait_or_while(until, timeout: int = 5) -> bool:
39
39
return False
40
40
41
41
42
+ def _raise_or_exit (val : bool , exit : bool ):
43
+ if exit :
44
+ raise Exit (int (not val ))
45
+ return val
46
+
47
+
42
48
def write_supervisor_config (cfg_json : str , _exit : Annotated [bool , Argument (hidden = True )] = True ):
43
49
cfg_obj = SupervisorAirflowConfiguration .model_validate_json (cfg_json )
44
50
if not _check_same (cfg_obj ):
45
51
log .critical ("Configs don't match" )
46
52
cfg_obj ._write_self ()
47
- return Exit ( 0 ) if _exit else True
53
+ return _raise_or_exit ( True , _exit )
48
54
49
55
50
56
def start_supervisor (
@@ -57,13 +63,13 @@ def start_supervisor(
57
63
if not _check_same (cfg_obj ):
58
64
log .critical ("Configs don't match" )
59
65
if _check_running (cfg_obj ):
60
- return Exit ( 0 ) if _exit else True
66
+ return _raise_or_exit ( True , _exit )
61
67
cfg_obj .start (daemon = True )
62
68
running = _wait_or_while (until = lambda : cfg_obj .running (), timeout = 30 )
63
69
if not running :
64
70
log .critical ("Still not running 30s after start command!" )
65
- return Exit ( 1 ) if _exit else False
66
- return Exit ( 0 ) if _exit else True
71
+ return _raise_or_exit ( False , _exit )
72
+ return _raise_or_exit ( True , _exit )
67
73
68
74
69
75
def start_programs (
@@ -77,21 +83,31 @@ def start_programs(
77
83
# TODO
78
84
ret = client .startAllProcesses ()
79
85
log .info (ret )
80
- return Exit ( 0 ) if _exit else True
86
+ return _raise_or_exit ( True , _exit )
81
87
82
88
83
89
def check_programs (
84
90
cfg : Annotated [
85
91
Path , Option (exists = True , file_okay = True , dir_okay = False , writable = False , readable = True , resolve_path = True )
86
92
],
93
+ check_running : bool = False ,
87
94
_exit : Annotated [bool , Argument (hidden = True )] = True ,
88
95
):
89
96
cfg_obj = SupervisorAirflowConfiguration .model_validate_json (cfg .read_text ())
90
97
client = SupervisorRemoteXMLRPCClient (cfg = cfg_obj )
91
98
# TODO
92
99
ret = client .getAllProcessInfo ()
93
- log .info (ret )
94
- return Exit (0 ) if _exit else True
100
+ for r in ret :
101
+ log .info (r .model_dump_json ())
102
+ if check_running :
103
+ meth = "running"
104
+ else :
105
+ meth = "ok"
106
+ if all (getattr (p , meth )() for p in ret ):
107
+ log .info ("all processes ok" )
108
+ return _raise_or_exit (True , _exit )
109
+ log .info ("processes not ok" )
110
+ return _raise_or_exit (False , _exit )
95
111
96
112
97
113
def restart_programs (
@@ -107,7 +123,7 @@ def restart_programs(
107
123
log .info (ret1 )
108
124
ret2 = client .startAllProcesses ()
109
125
log .info (ret2 )
110
- return Exit ( 0 ) if _exit else True
126
+ return _raise_or_exit ( True , _exit )
111
127
112
128
113
129
def stop_supervisor (
@@ -121,8 +137,8 @@ def stop_supervisor(
121
137
not_running = _wait_or_while (until = lambda : not cfg_obj .running (), timeout = 30 )
122
138
if not not_running :
123
139
log .critical ("Still running 30s after stop command!" )
124
- return Exit ( 1 ) if _exit else False
125
- return Exit ( 0 ) if _exit else True
140
+ return _raise_or_exit ( False , _exit )
141
+ return _raise_or_exit ( True , _exit )
126
142
127
143
128
144
def kill_supervisor (
@@ -136,8 +152,8 @@ def kill_supervisor(
136
152
still_running = _wait_or_while (until = lambda : not cfg_obj .running (), timeout = 30 )
137
153
if still_running :
138
154
log .critical ("Still running 30s after kill command!" )
139
- return Exit ( 1 ) if _exit else False
140
- return Exit ( 0 ) if _exit else True
155
+ return _raise_or_exit ( False , _exit )
156
+ return _raise_or_exit ( True , _exit )
141
157
142
158
143
159
def remove_supervisor_config (
@@ -152,14 +168,14 @@ def remove_supervisor_config(
152
168
still_running = kill_supervisor (cfg_obj , _exit = False )
153
169
154
170
if still_running :
155
- return Exit ( 1 ) if _exit else True
171
+ return _raise_or_exit ( False , _exit )
156
172
157
173
# TODO move to config
158
174
sleep (5 )
159
175
160
176
# TODO make optional
161
177
cfg_obj .rmdir ()
162
- return Exit ( 0 ) if _exit else True
178
+ return _raise_or_exit ( True , _exit )
163
179
164
180
165
181
def _add_to_typer (app , command : _SupervisorTaskStep , foo ):
0 commit comments