13
13
__all__ = ("Supervisor" ,)
14
14
15
15
16
- class Supervisor (DAG ):
17
- _supervisor_cfg : SupervisorAirflowConfiguration
18
- _supervisor_kill : DAG
19
- _supervisor_xmlrpc_client : SupervisorRemoteXMLRPCClient
16
+ class Supervisor (object ):
17
+ _dag : DAG
18
+ _cfg : SupervisorAirflowConfiguration
19
+ _kill_dag : DAG
20
+ _xmlrpc_client : SupervisorRemoteXMLRPCClient
20
21
21
- def __init__ (self , supervisor_cfg : SupervisorAirflowConfiguration , ** kwargs ):
22
+ def __init__ (self , dag : DAG , cfg : SupervisorAirflowConfiguration , ** kwargs ):
22
23
# store config
23
- self ._supervisor_cfg = supervisor_cfg
24
- self ._supervisor_xmlrpc_client = kwargs .pop (
25
- "supervisor_xmlrpc_client" , SupervisorRemoteXMLRPCClient (self ._supervisor_cfg )
26
- )
24
+ self ._cfg = cfg
27
25
28
- # setup role and tweak dag id
29
- if "dag_id" not in kwargs :
30
- kwargs ["dag_id" ] = list (self ._supervisor_cfg .program .keys ())[0 ]
26
+ # store or create client
27
+ self ._xmlrpc_client = kwargs .pop ("xmlrpc_client" , SupervisorRemoteXMLRPCClient (self ._cfg ))
31
28
32
- # override dag kwargs that dont make sense
33
- kwargs ["catchup" ] = False
34
- kwargs ["concurrency" ] = 1
35
- kwargs ["max_active_tasks" ] = 1
36
- kwargs ["max_active_runs" ] = 1
29
+ # store dag
30
+ self ._dag = dag
37
31
38
- # init with base DAG
39
- super ().__init__ (** kwargs )
32
+ self .setup_dag ()
40
33
41
34
# initialize tasks
42
35
self .initialize_tasks ()
@@ -52,6 +45,13 @@ def __init__(self, supervisor_cfg: SupervisorAirflowConfiguration, **kwargs):
52
45
# Default non running
53
46
PythonOperator (task_id = "skip" , python_callable = skip_ ) >> self ._force_kill
54
47
48
+ def setup_dag (self ):
49
+ # override dag kwargs that dont make sense
50
+ self ._dag .catchup = False
51
+ self ._dag .concurrency = 1
52
+ self ._dag .max_active_tasks = 1
53
+ self ._dag .max_active_runs = 1
54
+
55
55
def initialize_tasks (self ):
56
56
# tasks
57
57
self ._configure_supervisor = self .get_step_operator (step = "configure-supervisor" )
@@ -99,37 +99,35 @@ def unconfigure_supervisor(self) -> Operator:
99
99
100
100
@property
101
101
def supervisor_client (self ) -> SupervisorRemoteXMLRPCClient :
102
- return SupervisorRemoteXMLRPCClient (self ._supervisor_cfg )
102
+ return SupervisorRemoteXMLRPCClient (self ._cfg )
103
103
104
104
def get_base_operator_kwargs (self ) -> Dict :
105
- return dict (dag = self )
105
+ return dict (dag = self . _dag )
106
106
107
107
def get_step_kwargs (self , step : SupervisorTaskStep ) -> Dict :
108
108
if step == "configure-supervisor" :
109
109
from .commands import write_supervisor_config
110
110
111
- return dict (
112
- python_callable = lambda : write_supervisor_config (self ._supervisor_cfg , _exit = False ), do_xcom_push = True
113
- )
111
+ return dict (python_callable = lambda : write_supervisor_config (self ._cfg , _exit = False ), do_xcom_push = True )
114
112
elif step == "start-supervisor" :
115
113
from .commands import start_supervisor
116
114
117
115
return dict (
118
- python_callable = lambda : start_supervisor (self ._supervisor_cfg ._pydantic_path , _exit = False ),
116
+ python_callable = lambda : start_supervisor (self ._cfg ._pydantic_path , _exit = False ),
119
117
do_xcom_push = True ,
120
118
)
121
119
elif step == "start-programs" :
122
120
from .commands import start_programs
123
121
124
- return dict (python_callable = lambda : start_programs (self ._supervisor_cfg , _exit = False ), do_xcom_push = True )
122
+ return dict (python_callable = lambda : start_programs (self ._cfg , _exit = False ), do_xcom_push = True )
125
123
elif step == "stop-programs" :
126
124
from .commands import stop_programs
127
125
128
- return dict (python_callable = lambda : stop_programs (self ._supervisor_cfg , _exit = False ), do_xcom_push = True )
126
+ return dict (python_callable = lambda : stop_programs (self ._cfg , _exit = False ), do_xcom_push = True )
129
127
elif step == "check-programs" :
130
128
from .commands import check_programs
131
129
132
- def _check_programs (supervisor_cfg = self ._supervisor_cfg , ** kwargs ) -> CheckResult :
130
+ def _check_programs (supervisor_cfg = self ._cfg , ** kwargs ) -> CheckResult :
133
131
# TODO formalize
134
132
if check_programs (supervisor_cfg , check_done = True , _exit = False ):
135
133
# finish
@@ -144,35 +142,33 @@ def _check_programs(supervisor_cfg=self._supervisor_cfg, **kwargs) -> CheckResul
144
142
elif step == "restart-programs" :
145
143
from .commands import restart_programs
146
144
147
- return dict (python_callable = lambda : restart_programs (self ._supervisor_cfg , _exit = False ), do_xcom_push = True )
145
+ return dict (python_callable = lambda : restart_programs (self ._cfg , _exit = False ), do_xcom_push = True )
148
146
elif step == "stop-supervisor" :
149
147
from .commands import stop_supervisor
150
148
151
- return dict (python_callable = lambda : stop_supervisor (self ._supervisor_cfg , _exit = False ), do_xcom_push = True )
149
+ return dict (python_callable = lambda : stop_supervisor (self ._cfg , _exit = False ), do_xcom_push = True )
152
150
elif step == "unconfigure-supervisor" :
153
151
from .commands import remove_supervisor_config
154
152
155
- return dict (
156
- python_callable = lambda : remove_supervisor_config (self ._supervisor_cfg , _exit = False ), do_xcom_push = True
157
- )
153
+ return dict (python_callable = lambda : remove_supervisor_config (self ._cfg , _exit = False ), do_xcom_push = True )
158
154
elif step == "force-kill" :
159
155
from .commands import kill_supervisor
160
156
161
- return dict (python_callable = lambda : kill_supervisor (self ._supervisor_cfg , _exit = False ), do_xcom_push = True )
157
+ return dict (python_callable = lambda : kill_supervisor (self ._cfg , _exit = False ), do_xcom_push = True )
162
158
raise NotImplementedError
163
159
164
160
def get_step_operator (self , step : SupervisorTaskStep ) -> Operator :
165
161
if step == "check-programs" :
166
162
return HighAvailabilityOperator (
167
163
** {
168
- "task_id" : f"{ self .dag_id } -{ step } " ,
169
- "poke_interval" : self ._supervisor_cfg .check_interval .total_seconds (),
170
- "timeout" : self ._supervisor_cfg .check_timeout .total_seconds (),
164
+ "task_id" : f"{ self ._dag . dag_id } -{ step } " ,
165
+ "poke_interval" : self ._cfg .check_interval .total_seconds (),
166
+ "timeout" : self ._cfg .check_timeout .total_seconds (),
171
167
"mode" : "poke" ,
172
168
** self .get_base_operator_kwargs (),
173
169
** self .get_step_kwargs (step ),
174
170
}
175
171
)
176
172
return PythonOperator (
177
- ** {"task_id" : f"{ self .dag_id } -{ step } " , ** self .get_base_operator_kwargs (), ** self .get_step_kwargs (step )}
173
+ ** {"task_id" : f"{ self ._dag . dag_id } -{ step } " , ** self .get_base_operator_kwargs (), ** self .get_step_kwargs (step )}
178
174
)
0 commit comments