Discard semicolon stripping in SQL hook (#25855)

apache · Aug 21, 2022 · 874a95c · 874a95c
1 parent ccdd73e
commit 874a95c
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 17 deletions.
diff --git a/airflow/providers/common/sql/hooks/sql.py b/airflow/providers/common/sql/hooks/sql.py
@@ -244,9 +244,7 @@ def split_sql_string(sql: str) -> List[str]:
         :return: list of individual expressions
         """
         splits = sqlparse.split(sqlparse.format(sql, strip_comments=True))
-        statements: List[str] = list(
-            filter(None, [s.rstrip(';').strip() if s.endswith(';') else s.strip() for s in splits])
-        )
+        statements: List[str] = list(filter(None, splits))
         return statements
 
     def run(
@@ -278,7 +276,7 @@ def run(
             if split_statements:
                 sql = self.split_sql_string(sql)
             else:
-                sql = [self.strip_sql_string(sql)]
+                sql = [sql]
 
         if sql:
             self.log.debug("Executing following statements against DB: %s", list(sql))

diff --git a/tests/providers/common/sql/hooks/test_sqlparse.py b/tests/providers/common/sql/hooks/test_sqlparse.py
@@ -24,16 +24,13 @@
     "line,parsed_statements",
     [
         ('SELECT * FROM table', ['SELECT * FROM table']),
-        ('SELECT * FROM table;', ['SELECT * FROM table']),
-        ('SELECT * FROM table; # comment', ['SELECT * FROM table']),
-        ('SELECT * FROM table; # comment;', ['SELECT * FROM table']),
-        (' SELECT * FROM table ; # comment;', ['SELECT * FROM table']),
-        ('SELECT * FROM table;;;;;', ['SELECT * FROM table']),
-        ('SELECT * FROM table;;# comment;;;', ['SELECT * FROM table']),
-        ('SELECT * FROM table;;# comment;; ;', ['SELECT * FROM table']),
+        ('SELECT * FROM table;', ['SELECT * FROM table;']),
+        ('SELECT * FROM table; # comment', ['SELECT * FROM table;']),
+        ('SELECT * FROM table; # comment;', ['SELECT * FROM table;']),
+        (' SELECT * FROM table ; # comment;', ['SELECT * FROM table ;']),
         (
             'SELECT * FROM table; SELECT * FROM table2 # comment',
-            ['SELECT * FROM table', 'SELECT * FROM table2'],
+            ['SELECT * FROM table;', 'SELECT * FROM table2'],
         ),
     ],
 )

diff --git a/tests/providers/databricks/hooks/test_databricks_sql.py b/tests/providers/databricks/hooks/test_databricks_sql.py
@@ -83,7 +83,7 @@ def test_query(self, mock_requests, mock_conn):
         assert schema == test_schema
         assert results == []
 
-        cur.execute.assert_has_calls([mock.call(q) for q in [query.rstrip(';')]])
+        cur.execute.assert_has_calls([mock.call(q) for q in [query]])
         cur.close.assert_called()
 
     def test_no_query(self):

diff --git a/tests/providers/oracle/hooks/test_oracle.py b/tests/providers/oracle/hooks/test_oracle.py
@@ -268,7 +268,7 @@ def getvalue(self):
 
         self.cur.bindvars = None
         result = self.db_hook.callproc('proc', True, parameters)
-        assert self.cur.execute.mock_calls == [mock.call('BEGIN proc(); END')]
+        assert self.cur.execute.mock_calls == [mock.call('BEGIN proc(); END;')]
         assert result == parameters
 
     def test_callproc_dict(self):
@@ -280,7 +280,7 @@ def getvalue(self):
 
         self.cur.bindvars = {k: bindvar(v) for k, v in parameters.items()}
         result = self.db_hook.callproc('proc', True, parameters)
-        assert self.cur.execute.mock_calls == [mock.call('BEGIN proc(:a,:b,:c); END', parameters)]
+        assert self.cur.execute.mock_calls == [mock.call('BEGIN proc(:a,:b,:c); END;', parameters)]
         assert result == parameters
 
     def test_callproc_list(self):
@@ -292,7 +292,7 @@ def getvalue(self):
 
         self.cur.bindvars = list(map(bindvar, parameters))
         result = self.db_hook.callproc('proc', True, parameters)
-        assert self.cur.execute.mock_calls == [mock.call('BEGIN proc(:1,:2,:3); END', parameters)]
+        assert self.cur.execute.mock_calls == [mock.call('BEGIN proc(:1,:2,:3); END;', parameters)]
         assert result == parameters
 
     def test_callproc_out_param(self):
@@ -306,7 +306,7 @@ def bindvar(value):
         self.cur.bindvars = [bindvar(p() if type(p) is type else p) for p in parameters]
         result = self.db_hook.callproc('proc', True, parameters)
         expected = [1, 0, 0.0, False, '']
-        assert self.cur.execute.mock_calls == [mock.call('BEGIN proc(:1,:2,:3,:4,:5); END', expected)]
+        assert self.cur.execute.mock_calls == [mock.call('BEGIN proc(:1,:2,:3,:4,:5); END;', expected)]
         assert result == expected
 
     def test_test_connection_use_dual_table(self):