diff --git a/airflow/providers/common/sql/hooks/sql.py b/airflow/providers/common/sql/hooks/sql.py index cb26f73517a99..c94e16e0f25d6 100644 --- a/airflow/providers/common/sql/hooks/sql.py +++ b/airflow/providers/common/sql/hooks/sql.py @@ -244,7 +244,9 @@ def split_sql_string(sql: str) -> List[str]: :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 0000000000000..3137f3cd21942 --- /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