diff --git a/sqlparse/engine/statement_splitter.py b/sqlparse/engine/statement_splitter.py index bc57d170..e4591b07 100644 --- a/sqlparse/engine/statement_splitter.py +++ b/sqlparse/engine/statement_splitter.py @@ -21,6 +21,8 @@ def _reset(self): self._parenthesis_level = 0 self._unconfirmed_start = None self._is_create = False + self._is_create_block = False + self._expect_create_body = False self._seen_begin = False self.consume_ws = False @@ -98,7 +100,17 @@ def _change_splitlevel(self, ttype, value): elif ttype is T.Punctuation and value == ')': self._parenthesis_level = max(0, self._parenthesis_level - 1) return -1 - elif ttype not in T.Keyword: # if normal token return + + if (self._expect_create_body and ttype not in T.Whitespace + and ttype not in T.Comment): + self._expect_create_body = False + if ttype in T.Literal: + return 0 + if not (ttype is T.Keyword and value.upper() == 'BEGIN'): + self._block_stack.append('DECLARE') + return 1 + + if ttype not in T.Keyword: # if normal token return return 0 # Everything after here is ttype = T.Keyword @@ -109,6 +121,16 @@ def _change_splitlevel(self, ttype, value): self._is_create = True return 0 + if self._is_create and unified in ( + 'FUNCTION', 'PROCEDURE', 'PACKAGE', 'TRIGGER'): + self._is_create_block = True + return 0 + + if (self._is_create_block and unified in ('AS', 'IS') + and not self._block_stack): + self._expect_create_body = True + return 0 + # Handle DECLARE block start (only for CREATE statements) if unified == 'DECLARE' and self._is_create and not self._block_stack: self._block_stack.append('DECLARE') diff --git a/tests/test_split.py b/tests/test_split.py index 92c3fefe..cffb2d13 100644 --- a/tests/test_split.py +++ b/tests/test_split.py @@ -256,6 +256,34 @@ def test_split_begin_end_procedure(): # issue809 assert 'CREATE OR REPLACE PROCEDURE' in stmts[0] +def test_split_oracle_procedure_with_declarations(): # issue692 + sql = """CREATE PROCEDURE remove_emp (employee_id NUMBER) AS + tot_emps NUMBER; +BEGIN + DELETE FROM employees + WHERE employees.employee_id = remove_emp.employee_id; + tot_emps := tot_emps - 1; +END;""" + stmts = sqlparse.split(sql) + assert len(stmts) == 1 + assert 'tot_emps NUMBER;' in stmts[0] + + +def test_split_function_with_dollar_quoted_body_and_following_statement(): + sql = """CREATE FUNCTION foo() RETURNS integer AS +$body$ +BEGIN + select * from foo; +END; +$body$ +LANGUAGE 'plpgsql'; +SELECT 1;""" + stmts = sqlparse.split(sql) + assert len(stmts) == 2 + assert stmts[0].startswith('CREATE FUNCTION') + assert stmts[1] == 'SELECT 1;' + + def test_split_begin_transaction(): # issue826 # BEGIN TRANSACTION should not be treated as a block start sql = """BEGIN TRANSACTION; @@ -373,4 +401,3 @@ def test_split_standalone_for_update(): assert stmts[0] == "SELECT * FROM foo FOR UPDATE;" assert stmts[1] == "SELECT 3;" -