Merge pull request #30 from headout/ft-sql-triggerer

feat: moving db polling to triggers

Author: joelkbiju12

src/operators/deferred_job_result.py

@@ -7,6 +7,7 @@
 from airflow.triggers.temporal import TimeDeltaTrigger
 from ergo.exceptions import ErgoFailedResultException
 from ergo.models import ErgoJob, ErgoTask
+from ergo.triggers.task_poll import TaskPollTrigger
 from sqlalchemy.orm import joinedload
 from airflow.triggers.temporal import TimeDeltaTrigger
 
@@ -18,6 +19,7 @@ def __init__(
             self,
             pusher_task_id: str,
             wait_for_state=list(State.finished),
+            task_poll_trigger=True,
             *args,
             **kwargs
     ):
@@ -26,12 +28,14 @@ def __init__(
         kwargs.pop('ti_dict', None)
         super().__init__(*args, **kwargs)
         self.pusher_task_id = pusher_task_id
+        self.task_poll_trigger = task_poll_trigger
         if not isinstance(wait_for_state, (list, tuple)):
             wait_for_state = (wait_for_state,)
         self.wait_for_state = list(State.finished)
         if self.wait_for_state != wait_for_state:
             self.wait_for_state.extend(wait_for_state)
 
+
     def _get_ergo_task(self, ti_dict, session=None):
         return (
             session.query(ErgoTask)
@@ -40,17 +44,12 @@ def _get_ergo_task(self, ti_dict, session=None):
         ).one()
 
     @provide_session
-    def execute(self, context, session=None, event=None):
-        ti_dict = context.get('ti_dict', dict())
-        if not ti_dict:
-            ti = context['ti']
-            ti_dict['dag_id'] = ti.dag_id
-            ti_dict['run_id'] = ti.run_id
+    def _get_task_status(self, context, ti_dict, session=None):
         task = self._get_ergo_task(ti_dict, session=session)
         job = task.job
 
+        # Fallback from triggerer to worker polling if triggerer fails
         while task.state not in self.wait_for_state:
-            self.defer(trigger=TimeDeltaTrigger(timedelta(seconds=15)), method_name="execute")
             task = self._get_ergo_task(ti_dict, session=session)
             self.log.info('Received task - %s... STATE: %s', str(task), task.state)
             job = task.job
@@ -61,14 +60,42 @@ def execute(self, context, session=None, event=None):
                 self.log.info('Waiting for task "%s" to be queued...', str(task))
                 self.log.info('Waiting for task "%s" to reach state %s...', str(task), self.wait_for_state)
 
+            if self.task_poll_trigger:
+                self.defer(trigger=TimeDeltaTrigger(timedelta(seconds=20)), method_name="execute_complete")
+
         if task.state == State.FAILED:
             if job is not None:
                 self.log.info('Job - (%s)' + (f'responded back at {job.response_at}' if job.response_at else ''), str(job))
                 raise ErgoFailedResultException(job.result_code, job.error_msg)
             else:
                 raise ErgoFailedResultException(400, "Cron execution failed due to unknown reason")
 
+        self.xcom_push(context, "ergo_task_state", "success")
         self.log.info('Task - %s reached state %s', str(task), task.state)
+
+
+    def execute(self, context, event=None):
+        ti_dict = context.get('ti_dict', dict())
+        self.xcom_push(context, "ergo_task_state", "started")
+        if not ti_dict:
+            ti = context['ti']
+            ti_dict['dag_id'] = ti.dag_id
+            ti_dict['run_id'] = ti.run_id
+        if self.task_poll_trigger:
+            self.log.info('Polling DB task status using task poll trigger. Check triggerer logs to get more state info')
+            self.defer(trigger=TaskPollTrigger(ti_dict, self.pusher_task_id, self.wait_for_state, 20), method_name="execute_complete")
+        else:
+            self._get_task_status(context, ti_dict)
+        return
+
+    def execute_complete(self, context, event=None):
+        ti_dict = context.get('ti_dict', dict())
+        if not ti_dict:
+            ti = context['ti']
+            ti_dict['dag_id'] = ti.dag_id
+            ti_dict['run_id'] = ti.run_id
+        self.log.info("Control transferred from trigger to worker for remaining operator execution")
+        self._get_task_status(context, ti_dict)
         return
 
 

src/triggers/__init__.py

@@ -0,0 +1,72 @@
+import asyncio
+import os
+from concurrent.futures import ThreadPoolExecutor
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+from airflow.utils.db import provide_session
+from airflow.utils.state import State
+from ergo.exceptions import ErgoFailedResultException
+from ergo.models import ErgoJob, ErgoTask
+from sqlalchemy.orm import joinedload
+
+os.environ['PYTHONASYNCIODEBUG'] = '1'
+
+
+class TaskPollTrigger(BaseTrigger):
+
+    def __init__(
+            self,
+            ti_dict,
+            pusher_task_id: str,
+            wait_for_state=list(State.finished),
+            poke_interval: float = 20,
+    ):
+        super().__init__()
+        self.ti_dict = ti_dict
+        self.pusher_task_id = pusher_task_id
+        self.wait_for_state = wait_for_state
+        self.poke_interval = poke_interval
+
+    def serialize(self):
+        return (
+            "ergo.triggers.task_poll.TaskPollTrigger",
+            {
+                "ti_dict": self.ti_dict,
+                "pusher_task_id": self.pusher_task_id,
+                "wait_for_state": self.wait_for_state,
+                "poke_interval": self.poke_interval,
+            },
+        )
+
+    async def _get_ergo_task(self, session=None):
+        return (
+            session.query(ErgoTask)
+            .options(joinedload('job'))
+            .filter_by(ti_task_id=self.pusher_task_id, ti_dag_id=self.ti_dict['dag_id'], ti_run_id=self.ti_dict['run_id'])
+        ).one()
+
+
+    @provide_session
+    async def _check_task_status(self, session=None):
+        task = await self._get_ergo_task(session=session)
+        job = task.job
+
+        if task.state not in self.wait_for_state:
+            self.log.info('Received task - %s... STATE: %s', str(task), task.state)
+            job = task.job
+            if job is not None:
+                self.log.info(
+                    'Job - (%s)' + (f'responded back at {job.response_at}' if job.response_at else ''), str(job))
+            else:
+                self.log.info('Waiting for task "%s" to be queued...', str(task))
+                self.log.info('Waiting for task "%s" to reach state %s...', str(task), self.wait_for_state)
+            return False
+
+        self.log.info('Task - %s reached state %s', str(task), task.state)
+        return True
+
+    async def run(self):
+        while True:
+            task_completed = await self._check_task_status()
+            if task_completed:
+                yield TriggerEvent(True)
+            await asyncio.sleep(self.poke_interval)