Keep Oracle procedure declarations in split statement

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')

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;"
 
-