[airflow] branch main updated: Fix SQL split string to include `;-less` statements (#25713)

[potiuk]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 7a19651369 Fix SQL split string to include `;-less` statements (#25713)
7a19651369 is described below

commit 7a19651369790e2abb563d96a42f41ec31ebfb85
Author: Jarek Potiuk <jarek.pot...@polidea.com>
AuthorDate: Mon Aug 15 11:28:04 2022 +0200

    Fix SQL split string to include `;-less` statements (#25713)
    
    There was a bug in an incoming change to common-sql provider
    introduced in #23971 where `;-less` statements were removed
    when "split_statements" flag was used. Since this flag is used
    by default in Databricks statement, it introduced backwards
    incompatible change.
---
 airflow/providers/common/sql/hooks/sql.py         |  4 ++-
 tests/providers/common/sql/hooks/test_sqlparse.py | 41 +++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/common/sql/hooks/sql.py 
b/airflow/providers/common/sql/hooks/sql.py
index cb26f73517..c94e16e0f2 100644
--- a/airflow/providers/common/sql/hooks/sql.py
+++ b/airflow/providers/common/sql/hooks/sql.py
@@ -244,7 +244,9 @@ class DbApiHook(BaseForDbApiHook):
         :return: list of individual expressions
         """
         splits = sqlparse.split(sqlparse.format(sql, strip_comments=True))
-        statements = [s.rstrip(';') for s in splits if s.endswith(';')]
+        statements: List[str] = list(
+            filter(None, [s.rstrip(';').strip() if s.endswith(';') else 
s.strip() for s in splits])
+        )
         return statements
 
     def run(
diff --git a/tests/providers/common/sql/hooks/test_sqlparse.py 
b/tests/providers/common/sql/hooks/test_sqlparse.py
new file mode 100644
index 0000000000..3137f3cd21
--- /dev/null
+++ b/tests/providers/common/sql/hooks/test_sqlparse.py
@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+
+@pytest.mark.parametrize(
+    "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 table2 # comment',
+            ['SELECT * FROM table', 'SELECT * FROM table2'],
+        ),
+    ],
+)
+def test_sqlparse(line, parsed_statements):
+    assert DbApiHook.split_sql_string(line) == parsed_statements