Skip to content

Commit

Permalink
Fix SQL split string to include ;-less statements (#25713)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
potiuk committed Aug 15, 2022
1 parent 8829283 commit 7a19651
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 1 deletion.
4 changes: 3 additions & 1 deletion airflow/providers/common/sql/hooks/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
41 changes: 41 additions & 0 deletions tests/providers/common/sql/hooks/test_sqlparse.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 7a19651

Please sign in to comment.