Refactor statement splitter to a simplified stack-based architecture (fixes #845)

Author: andialbrecht

CHANGELOG

@@ -1,7 +1,9 @@
 Development Version
 -------------------
 
-Nothing yet.
+Bug Fixes
+
+* Fix statement splitting (issue845).
 
 
 Release 0.5.5 (Dec 19, 2025)

sqlparse/engine/statement_splitter.py

@@ -16,10 +16,10 @@ def __init__(self):
 
     def _reset(self):
         """Set the filter attributes to its default values"""
-        self._in_declare = False
-        self._in_case = False
+        self._block_stack = []
+        self._parenthesis_level = 0
+        self._unconfirmed_start = None
         self._is_create = False
-        self._begin_depth = 0
         self._seen_begin = False
 
         self.consume_ws = False
@@ -29,37 +29,44 @@ def _reset(self):
     def _change_splitlevel(self, ttype, value):
         """Get the new split level (increase, decrease or remain equal)"""
 
+        # Semicolon resets unconfirmed loop starters
+        if ttype is T.Punctuation and value == ';':
+            self._unconfirmed_start = None
+
         # parenthesis increase/decrease a level
         if ttype is T.Punctuation and value == '(':
+            self._parenthesis_level += 1
             return 1
         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
             return 0
 
         # Everything after here is ttype = T.Keyword
-        # Also to note, once entered an If statement you are done and basically
-        # returning
         unified = value.upper()
 
-        # three keywords begin with CREATE, but only one of them is DDL
         # DDL Create though can contain more words such as "or replace"
         if ttype is T.Keyword.DDL and unified.startswith('CREATE'):
             self._is_create = True
             return 0
 
-        # can have nested declare inside of being...
-        if unified == 'DECLARE' and self._is_create and self._begin_depth == 0:
-            self._in_declare = True
+        # 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')
             return 1
 
+        # Handle BEGIN block start
         if unified == 'BEGIN':
-            self._begin_depth += 1
             self._seen_begin = True
-            if self._is_create:
-                # FIXME(andi): This makes no sense.  ## this comment neither
+            # Transition DECLARE to BEGIN if present
+            if self._block_stack and self._block_stack[-1] == 'DECLARE':
+                self._block_stack.pop()
+                self._block_stack.append('BEGIN')
+                return 0
+            else:
+                self._block_stack.append('BEGIN')
                 return 1
-            return 0
 
         # Issue826: If we see a transaction keyword after BEGIN,
         # it's a transaction statement, not a block.
@@ -68,28 +75,72 @@ def _change_splitlevel(self, ttype, value):
                 unified in ('TRANSACTION', 'WORK', 'TRAN',
                             'DISTRIBUTED', 'DEFERRED',
                             'IMMEDIATE', 'EXCLUSIVE'):
-            self._begin_depth = max(0, self._begin_depth - 1)
             self._seen_begin = False
+            if self._block_stack and self._block_stack[-1] == 'BEGIN':
+                self._block_stack.pop()
+                return -1
             return 0
 
-        # BEGIN and CASE/WHEN both end with END
-        if unified == 'END':
-            if not self._in_case:
-                self._begin_depth = max(0, self._begin_depth - 1)
-            else:
-                self._in_case = False
-            return -1
-
-        if (unified in ('IF', 'FOR', 'WHILE', 'CASE')
-                and self._is_create and self._begin_depth > 0):
-            if unified == 'CASE':
-                self._in_case = True
-            return 1
+        # Inside a block, check for nested loop or control structures
+        if 'BEGIN' in self._block_stack:
+            if unified == 'FOR':
+                self._unconfirmed_start = 'FOR'
+                return 0
+            elif unified == 'WHILE':
+                self._unconfirmed_start = 'WHILE'
+                return 0
+            elif unified == 'LOOP':
+                if self._unconfirmed_start in ('FOR', 'WHILE'):
+                    self._block_stack.append(self._unconfirmed_start)
+                    self._unconfirmed_start = None
+                    return 1
+                else:
+                    self._block_stack.append('LOOP')
+                    return 1
+            elif unified == 'DO':
+                if self._unconfirmed_start in ('FOR', 'WHILE'):
+                    self._block_stack.append(self._unconfirmed_start)
+                    self._unconfirmed_start = None
+                    return 1
+            elif unified == 'IF':
+                self._block_stack.append('IF')
+                return 1
+            elif unified == 'CASE':
+                self._block_stack.append('CASE')
+                return 1
 
-        if unified in ('END IF', 'END FOR', 'END WHILE'):
-            return -1
+        # Handle closing keywords
+        if unified == 'END IF':
+            if self._block_stack and self._block_stack[-1] == 'IF':
+                self._block_stack.pop()
+                return -1
+        elif unified == 'END FOR':
+            if self._block_stack and self._block_stack[-1] == 'FOR':
+                self._block_stack.pop()
+                return -1
+        elif unified == 'END WHILE':
+            if self._block_stack and self._block_stack[-1] == 'WHILE':
+                self._block_stack.pop()
+                return -1
+        elif unified == 'END LOOP':
+            if self._block_stack and self._block_stack[-1] in ('LOOP', 'FOR', 'WHILE'):
+                self._block_stack.pop()
+                return -1
+        elif unified == 'END CASE':
+            if self._block_stack and self._block_stack[-1] == 'CASE':
+                self._block_stack.pop()
+                return -1
+        elif unified == 'END':
+            if self._block_stack:
+                if self._block_stack[-1] in ('CASE', 'BEGIN'):
+                    self._block_stack.pop()
+                    return -1
+                else:
+                    self._block_stack.pop()
+                    return -1
+            else:
+                return -1
 
-        # Default
         return 0
 
     def process(self, stream):
@@ -125,10 +176,12 @@ def process(self, stream):
                 # If we just saw BEGIN; then this is a transaction BEGIN,
                 # not a BEGIN...END block, so decrement depth
                 if self._seen_begin:
-                    self._begin_depth = max(0, self._begin_depth - 1)
+                    if self._block_stack and self._block_stack[-1] == 'BEGIN':
+                        self._block_stack.pop()
+                        self.level = max(0, self.level - 1)
                 self._seen_begin = False
                 # Split on semicolon if not inside a BEGIN...END block
-                if self.level <= 0 and self._begin_depth == 0:
+                if self.level <= 0 and 'BEGIN' not in self._block_stack:
                     self.consume_ws = True
             elif ttype is T.Keyword and value.split()[0] == 'GO':
                 self.consume_ws = True

sqlparse/keywords.py

@@ -69,7 +69,7 @@
     (r'(?<![\w\])])(\[[^\]\[]+\])', tokens.Name),
     (r'((LEFT\s+|RIGHT\s+|FULL\s+)?(INNER\s+|OUTER\s+|STRAIGHT\s+)?'
      r'|(CROSS\s+|NATURAL\s+)?)?JOIN\b', tokens.Keyword),
-    (r'END(\s+IF|\s+LOOP|\s+WHILE)?\b', tokens.Keyword),
+    (r'END(\s+IF|\s+LOOP|\s+WHILE|\s+FOR|\s+CASE)?\b', tokens.Keyword),
     (r'IF\s+(NOT\s+)?EXISTS\b', tokens.Keyword),
     (r'NOT\s+NULL\b', tokens.Keyword),
     (r'(ASC|DESC)(\s+NULLS\s+(FIRST|LAST))?\b', tokens.Keyword.Order),

tests/test_split.py

@@ -284,3 +284,93 @@ def test_split_begin_transaction_formatted():  # issue826
     assert stmts[1].startswith('DELETE')
     assert stmts[2].startswith('INSERT')
     assert stmts[3] == 'END\nTRANSACTION;'
+
+
+def test_split_anonymous_begin_end_for():  # issue845 Case 1
+    sql = """
+BEGIN
+    SELECT 1;
+    FOR R DO
+        SELECT 1;
+    END FOR;
+END;
+"""
+    stmts = sqlparse.split(sql)
+    assert len(stmts) == 1
+    assert "END FOR;" in stmts[0]
+
+
+def test_split_anonymous_begin_end_case_inline():  # issue845 Case 2
+    sql = """
+BEGIN
+    SELECT 1;
+    IF 1 THEN
+        SELECT CASE WHEN 1 THEN 2 ELSE 3 END AS COUNT;
+    ELSE
+        SELECT 2;
+    END IF;
+END;
+"""
+    stmts = sqlparse.split(sql)
+    assert len(stmts) == 1
+    assert "END AS COUNT;" in stmts[0]
+
+
+def test_split_for_update_in_begin_end():
+    # Verify that FOR UPDATE / FOR SHARE inside a BEGIN ... END block do not break level balancing
+    sql = """
+BEGIN
+    SELECT * FROM foo FOR UPDATE;
+    SELECT * FROM bar FOR SHARE;
+END;
+SELECT 3;
+"""
+    stmts = sqlparse.split(sql)
+    assert len(stmts) == 2
+    assert "SELECT 3;" in stmts[1]
+
+
+def test_split_multiple_for_loops_in_begin_end():
+    # Verify that multiple sequential loops inside a BEGIN ... END block balance correctly
+    sql = """
+BEGIN
+    FOR x IN select_query LOOP
+        SELECT 1;
+    END LOOP;
+    FOR y IN select_query LOOP
+        SELECT 2;
+    END LOOP;
+END;
+SELECT 3;
+"""
+    stmts = sqlparse.split(sql)
+    assert len(stmts) == 2
+    assert "SELECT 3;" in stmts[1]
+
+
+def test_split_procedural_case_end_case():
+    # Verify that CASE closed by END CASE inside a BEGIN block balances correctly
+    sql = """
+BEGIN
+    CASE val
+        WHEN 1 THEN SELECT 'one';
+        WHEN 2 THEN SELECT 'two';
+        ELSE SELECT 'other';
+    END CASE;
+END;
+SELECT 3;
+"""
+    stmts = sqlparse.split(sql)
+    assert len(stmts) == 2
+    assert "SELECT 3;" in stmts[1]
+
+
+def test_split_standalone_for_update():
+    # Verify that standalone FOR UPDATE statements split correctly
+    sql = "SELECT * FROM foo FOR UPDATE; SELECT 3;"
+    stmts = sqlparse.split(sql)
+    assert len(stmts) == 2
+    assert stmts[0] == "SELECT * FROM foo FOR UPDATE;"
+    assert stmts[1] == "SELECT 3;"
+
+