def handle_event(self, event: Event, ti: TaskInstance, ts: TaskState, session=None)->TaskAction:
if ts.task_state is None:
ts.task_state = AiFlowTs()
af_ts = copy.deepcopy(ts.task_state)
af_ts.event_map[(event.key, event.event_type)] = event
aw = ActionWrapper()
res = self.met_sc(af_ts, aw)
if res:
if aw.action in SCHEDULED_ACTION:
af_ts.schedule_time = time.time_ns()
ts.task_state = af_ts
session.merge(ts)
session.commit()
if len(self.configs) == 0:
return TaskAction.START
else:
return aw.action
else:
ts.task_state = af_ts
session.merge(ts)
session.commit()
return TaskAction.NONE
def _operator_handle_event(event, operator,
execution_date) -> SchedulingAction:
task_state = TaskState.get_task_state(operator.dag_id,
operator.task_id, execution_date)
event_handler = operator.get_events_handler()
if not task_state:
task_state = TaskState(task_id=operator.task_id,
dag_id=operator.dag_id,
execution_date=execution_date)
with create_session() as session:
session.add(task_state)
session.commit()
scheduling_action, state = event_handler.handle_event(
event, task_state.task_state)
task_state.task_state = state
task_state.update_task_state()
return scheduling_action
