Enable semicolon stripping for PrestoHook (#34828)

apache · Aug 31, 2024 · 06278ce · 06278ce
1 parent 0e97fa8
commit 06278ce
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/airflow/providers/presto/hooks/presto.py b/airflow/providers/presto/hooks/presto.py
@@ -84,6 +84,7 @@ class PrestoHook(DbApiHook):
  default_conn_name = "presto_default"
  conn_type = "presto"
  hook_name = "Presto"
+ strip_semicolon = True
 
  def __init__(self, *args, **kwargs):
  super().__init__(*args, **kwargs)

diff --git a/tests/providers/presto/hooks/test_presto.py b/tests/providers/presto/hooks/test_presto.py
@@ -276,6 +276,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)
+
  def test_serialize_cell(self):
  assert "foo" == self.db_hook._serialize_cell("foo", None)
  assert 1 == self.db_hook._serialize_cell(1, None)