fix: listen state to receive sigterm cancellation (#17)

* fix: listen state to receive sigterm cancellation

* fix: no need to use thread for terminating jobs

* refactor: handle_job_cancel to register_job

* refactor: on_job_cancel to o_job_regsiter

task/bq2bq/executor/bumblebee/bigquery_service.py

@@ -53,7 +53,7 @@ def get_table(self, full_table_name):
 
 class BigqueryService(BaseBigqueryService):
 
-    def __init__(self, client, labels, writer, on_job_finish = None, on_job_cancelled = None):
+    def __init__(self, client, labels, writer, on_job_finish = None, on_job_register = None):
         """
 
         :rtype:
@@ -62,7 +62,7 @@ def __init__(self, client, labels, writer, on_job_finish = None, on_job_cancelle
         self.labels = labels
         self.writer = writer
         self.on_job_finish = on_job_finish
-        self.on_job_cancelled = on_job_cancelled
+        self.on_job_register = on_job_register
 
     def execute_query(self, query):
         query_job_config = QueryJobConfig()
@@ -78,8 +78,8 @@ def execute_query(self, query):
         logger.info("Job {} is initially in state {} of {} project".format(query_job.job_id, query_job.state,
                                                                            query_job.project))
 
-        if self.on_job_cancelled:
-            self.on_job_cancelled(self.client, query_job)
+        if self.on_job_register:
+            self.on_job_register(self.client, query_job)
 
         try:
             result = query_job.result()
@@ -129,8 +129,8 @@ def transform_load(self,
         logger.info("Job {} is initially in state {} of {} project".format(query_job.job_id, query_job.state,
                                                                            query_job.project))
 
-        if self.on_job_cancelled:
-            self.on_job_cancelled(self.client, query_job)
+        if self.on_job_register:
+            self.on_job_register(self.client, query_job)
 
         try:
             result = query_job.result()
@@ -174,7 +174,7 @@ def get_table(self, full_table_name):
         return self.client.get_table(table_ref)
 
 
-def create_bigquery_service(task_config: TaskConfigFromEnv, labels, writer, on_job_finish = None, on_job_cancelled = None):
+def create_bigquery_service(task_config: TaskConfigFromEnv, labels, writer, on_job_finish = None, on_job_register = None):
     if writer is None:
         writer = writer.StdWriter()
 
@@ -183,7 +183,7 @@ def create_bigquery_service(task_config: TaskConfigFromEnv, labels, writer, on_j
     default_query_job_config.priority = task_config.query_priority
     default_query_job_config.allow_field_addition = task_config.allow_field_addition
     client = bigquery.Client(project=task_config.execution_project, credentials=credentials, default_query_job_config=default_query_job_config)
-    return BigqueryService(client, labels, writer, on_job_finish=on_job_finish, on_job_cancelled=on_job_cancelled)
+    return BigqueryService(client, labels, writer, on_job_finish=on_job_finish, on_job_register=on_job_register)
 
 
 def _get_bigquery_credentials():

task/bq2bq/executor/bumblebee/bq2bq.py

@@ -22,7 +22,7 @@ def bq2bq(properties_file: str,
           labels: dict = {},
           output_on: str = './return.json',
           on_job_finish = None,
-          on_job_cancelled = None,
+          on_job_register = None,
           ):
 
     logger.info("Using bumblebee version: {}".format(VERSION))
@@ -39,7 +39,7 @@ def bq2bq(properties_file: str,
 
     bigquery_service = DummyService()
     if not dry_run:
-        bigquery_service = create_bigquery_service(task_config, job_labels, writer, on_job_finish=on_job_finish, on_job_cancelled=on_job_cancelled)
+        bigquery_service = create_bigquery_service(task_config, job_labels, writer, on_job_finish=on_job_finish, on_job_register=on_job_register)
 
     transformation = Transformation(bigquery_service,
                                     task_config,

task/bq2bq/executor/bumblebee/handler.py

@@ -8,20 +8,30 @@ class BigqueryJobHandler:
     def __init__(self) -> None:
         self._sum_slot_millis = 0
         self._sum_total_bytes_processed = 0
+        self.client = None
+        self.jobs = []
+        self._init_signal_handling()
+
+    def _init_signal_handling(self):
+        def handle_sigterm(signum, frame):
+            self._terminate_jobs()
+            sys.exit(1)
+        signal.signal(signal.SIGTERM, handle_sigterm)
+
+    def _terminate_jobs(self):
+        if self.client and self.jobs:
+            for job in self.jobs:
+                job_id = job.job_id
+                self.client.cancel_job(job_id)
+                logger.info(f"{job_id} successfully cancelled")
 
     def handle_job_finish(self, job) -> None:
         self._sum_slot_millis += job.slot_millis
         self._sum_total_bytes_processed += job.total_bytes_processed
 
-    def handle_job_cancelled(self, client, job):
-        c = client
-        job_id = job.job_id
-        def handler(signum, frame):
-            c.cancel_job(job_id)
-            logger.info(f"{job_id} successfully cancelled")
-            sys.exit(1)
-
-        signal.signal(signal.SIGTERM, handler)
+    def register_job(self, client, job):
+        self.client = client
+        self.jobs.append(job)
 
     def get_sum_slot_millis(self) -> int:
         return self._sum_slot_millis

task/bq2bq/executor/main.py

@@ -26,7 +26,7 @@
         app_config.job_labels,
         app_config.xcom_path,
         on_job_finish = job_handler.handle_job_finish,
-        on_job_cancelled = job_handler.handle_job_cancelled,
+        on_job_register = job_handler.register_job,
     )
 
     xcom_data['monitoring'] = {