Add on_kill() to kill Trino query if the task is killed (#24559)

This PR is a follow up to PR #24415.
It adds on_kill method to the TrinoOperator to kill Trino query if Airflow task is killed

airflow/providers/trino/hooks/trino.py

@@ -93,6 +93,7 @@ class TrinoHook(DbApiHook):
     default_conn_name = 'trino_default'
     conn_type = 'trino'
     hook_name = 'Trino'
+    query_id = ''
 
     def get_conn(self) -> Connection:
         """Returns a connection object"""
@@ -301,6 +302,7 @@ def run(
                 results = []
                 for sql_statement in sql:
                     self._run_command(cur, self._strip_sql(sql_statement), parameters)
+                    self.query_id = cur.stats["queryId"]
                     if handler is not None:
                         result = handler(cur)
                         results.append(result)

airflow/providers/trino/operators/trino.py

@@ -19,6 +19,8 @@
 
 from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence, Union
 
+from trino.exceptions import TrinoQueryError
+
 from airflow.models import BaseOperator
 from airflow.providers.trino.hooks.trino import TrinoHook
 
@@ -79,3 +81,18 @@ def execute(self, context: 'Context') -> None:
         self.hook.run(
             sql=self.sql, autocommit=self.autocommit, parameters=self.parameters, handler=self.handler
         )
+
+    def on_kill(self) -> None:
+        if self.hook is not None and isinstance(self.hook, TrinoHook):
+            query_id = "'" + self.hook.query_id + "'"
+            try:
+                self.log.info("Stopping query run with queryId - %s", self.hook.query_id)
+                self.hook.run(
+                    sql=f"CALL system.runtime.kill_query(query_id => {query_id},message => 'Job "
+                    f"killed by "
+                    f"user');",
+                    handler=list,
+                )
+            except TrinoQueryError as e:
+                self.log.info(str(e))
+            self.log.info("Trino query (%s) terminated", query_id)

tests/system/providers/trino/example_trino.py

@@ -45,7 +45,7 @@
     )
     trino_create_table = TrinoOperator(
         task_id="trino_create_table",
-        sql=f"""CREATE TABLE {SCHEMA}.{TABLE}(
+        sql=f"""CREATE TABLE IF NOT EXISTS {SCHEMA}.{TABLE}(
         cityid bigint,
         cityname varchar
         )""",
@@ -60,9 +60,9 @@
 
     trino_multiple_queries = TrinoOperator(
         task_id="trino_multiple_queries",
-        sql=f"""CREATE TABLE {SCHEMA}.{TABLE1}(cityid bigint,cityname varchar);
+        sql=f"""CREATE TABLE IF NOT EXISTS {SCHEMA}.{TABLE1}(cityid bigint,cityname varchar);
         INSERT INTO {SCHEMA}.{TABLE1} VALUES (2, 'San Jose');
-        CREATE TABLE {SCHEMA}.{TABLE2}(cityid bigint,cityname varchar);
+        CREATE TABLE IF NOT EXISTS {SCHEMA}.{TABLE2}(cityid bigint,cityname varchar);
         INSERT INTO {SCHEMA}.{TABLE2} VALUES (3, 'San Diego');""",
         handler=list,
     )