diff --git a/sqlparse/engine/statement_splitter.py b/sqlparse/engine/statement_splitter.py index bc57d170..47c9def5 100644 --- a/sqlparse/engine/statement_splitter.py +++ b/sqlparse/engine/statement_splitter.py @@ -150,15 +150,18 @@ def _change_splitlevel(self, ttype, value): def process(self, stream): """Process the stream""" - EOS_TTYPE = T.Whitespace, T.Comment.Single - # Run over all stream tokens for ttype, value in stream: - # Yield token if we finished a statement and there's no whitespaces - # It will count newline token as a non whitespace. In this context - # whitespace ignores newlines. - # why don't multi line comments also count? - if self.consume_ws and ttype not in EOS_TTYPE: + # Once a statement is finished (consume_ws), keep trailing + # whitespace and single-line comments with it, and only cut a + # new statement when real content arrives. Newlines are + # whitespace too, so a terminating semicolon's trailing newline + # stays with its statement instead of being split off into a + # whitespace-only remainder that gets dropped at end of stream + # (which broke str(parse(sql)) == sql for inputs ending in + # ";\n", the common end-of-file case). + if self.consume_ws and not ( + ttype in T.Whitespace or ttype in T.Comment.Single): yield sql.Statement(self.tokens) # Reset filter and prepare to process next statement diff --git a/tests/test_split.py b/tests/test_split.py index 92c3fefe..3c488625 100644 --- a/tests/test_split.py +++ b/tests/test_split.py @@ -152,6 +152,25 @@ def test_split_ignores_empty_newlines(): assert stmts[1] == 'select bar;' +@pytest.mark.parametrize('sql', [ + 'select 1;\n', + 'select 1;\n\n', + 'select 1;\r\n', + 'select 1; ', + 'select 1;\t', + 'select 1;\nselect 2;\n', + 'create table t (a int);\n\n', +]) +def test_split_preserves_trailing_whitespace(sql): + # The whitespace after a terminating semicolon must stay with the + # statement so that the parsed statements round-trip back to the input. + # A trailing newline (the common end-of-file "...;\n" case) used to be + # split off into a whitespace-only remainder and dropped, unlike a + # trailing space or tab. + stmts = sqlparse.parse(sql) + assert ''.join(str(q) for q in stmts) == sql + + def test_split_quotes_with_new_line(): stmts = sqlparse.split('select "foo\nbar"') assert len(stmts) == 1