Enable semicolon stripping for TrinoHook (#34828)

apache · Aug 31, 2024 · 580458f · 580458f
1 parent 06278ce
commit 580458f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/airflow/providers/trino/hooks/trino.py b/airflow/providers/trino/hooks/trino.py
@@ -86,6 +86,7 @@ class TrinoHook(DbApiHook):
  default_conn_name = "trino_default"
  conn_type = "trino"
  hook_name = "Trino"
+ strip_semicolon = True
  query_id = ""
  _test_connection_sql = "select 1"
 

diff --git a/tests/providers/trino/hooks/test_trino.py b/tests/providers/trino/hooks/test_trino.py
@@ -345,6 +345,13 @@ def test_get_pandas_df(self):
 
  self.cur.execute.assert_called_once_with(statement, None)
 
+ def test_split_sql_string(self):
+ statement = "SELECT 1; SELECT 2"
+ result_sets = ["SELECT 1", "SELECT 2"]
+ self.cur.fetchall.return_value = result_sets
+
+ assert result_sets == self.db_hook.split_sql_string(statement)
+
  @patch("airflow.providers.trino.hooks.trino.TrinoHook.run")
  def test_run(self, mock_run):
  sql = "SELECT 1"