Skip to content

Commit ec28d0e

Browse files
andialbrechtandialbrecht
committed
Refactor StatementSplitter to decompose split level logic and centralize standalone BEGIN; handling
1 parent 8f978ca commit ec28d0e

1 file changed

Lines changed: 61 additions & 63 deletions

File tree

‎sqlparse/engine/statement_splitter.py‎

Lines changed: 61 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,69 @@ def _reset(self):
2626
self.tokens = []
2727
self.level = 0
2828

29+
def _handle_nested_block(self, unified):
30+
"""Check for nested loop or control structures inside a block"""
31+
if unified == 'FOR':
32+
self._unconfirmed_start = 'FOR'
33+
return 0
34+
if unified == 'WHILE':
35+
self._unconfirmed_start = 'WHILE'
36+
return 0
37+
if unified in ('LOOP', 'DO'):
38+
if self._unconfirmed_start in ('FOR', 'WHILE'):
39+
self._block_stack.append(self._unconfirmed_start)
40+
self._unconfirmed_start = None
41+
return 1
42+
if unified == 'LOOP':
43+
self._block_stack.append('LOOP')
44+
return 1
45+
if unified in ('IF', 'CASE'):
46+
self._block_stack.append(unified)
47+
return 1
48+
return None
49+
50+
def _handle_closing_keyword(self, unified):
51+
"""Handle closing keywords for blocks"""
52+
if unified == 'END IF':
53+
if self._block_stack and self._block_stack[-1] == 'IF':
54+
self._block_stack.pop()
55+
return -1
56+
elif unified == 'END FOR':
57+
if self._block_stack and self._block_stack[-1] == 'FOR':
58+
self._block_stack.pop()
59+
return -1
60+
elif unified == 'END WHILE':
61+
if self._block_stack and self._block_stack[-1] == 'WHILE':
62+
self._block_stack.pop()
63+
return -1
64+
elif unified == 'END LOOP':
65+
if (self._block_stack and
66+
self._block_stack[-1] in ('LOOP', 'FOR', 'WHILE')):
67+
self._block_stack.pop()
68+
return -1
69+
elif unified == 'END CASE':
70+
if self._block_stack and self._block_stack[-1] == 'CASE':
71+
self._block_stack.pop()
72+
return -1
73+
elif unified == 'END':
74+
if self._block_stack:
75+
self._block_stack.pop()
76+
return -1
77+
return 0
78+
2979
def _change_splitlevel(self, ttype, value):
3080
"""Get the new split level (increase, decrease or remain equal)"""
3181

3282
# Semicolon resets unconfirmed loop starters
83+
# and handles standalone BEGIN;
3384
if ttype is T.Punctuation and value == ';':
3485
self._unconfirmed_start = None
86+
if self._seen_begin:
87+
self._seen_begin = False
88+
if self._block_stack and self._block_stack[-1] == 'BEGIN':
89+
self._block_stack.pop()
90+
return -1
91+
return 0
3592

3693
# parenthesis increase/decrease a level
3794
if ttype is T.Punctuation and value == '(':
@@ -83,65 +140,12 @@ def _change_splitlevel(self, ttype, value):
83140

84141
# Inside a block, check for nested loop or control structures
85142
if 'BEGIN' in self._block_stack:
86-
if unified == 'FOR':
87-
self._unconfirmed_start = 'FOR'
88-
return 0
89-
elif unified == 'WHILE':
90-
self._unconfirmed_start = 'WHILE'
91-
return 0
92-
elif unified == 'LOOP':
93-
if self._unconfirmed_start in ('FOR', 'WHILE'):
94-
self._block_stack.append(self._unconfirmed_start)
95-
self._unconfirmed_start = None
96-
return 1
97-
else:
98-
self._block_stack.append('LOOP')
99-
return 1
100-
elif unified == 'DO':
101-
if self._unconfirmed_start in ('FOR', 'WHILE'):
102-
self._block_stack.append(self._unconfirmed_start)
103-
self._unconfirmed_start = None
104-
return 1
105-
elif unified == 'IF':
106-
self._block_stack.append('IF')
107-
return 1
108-
elif unified == 'CASE':
109-
self._block_stack.append('CASE')
110-
return 1
143+
res = self._handle_nested_block(unified)
144+
if res is not None:
145+
return res
111146

112147
# Handle closing keywords
113-
if unified == 'END IF':
114-
if self._block_stack and self._block_stack[-1] == 'IF':
115-
self._block_stack.pop()
116-
return -1
117-
elif unified == 'END FOR':
118-
if self._block_stack and self._block_stack[-1] == 'FOR':
119-
self._block_stack.pop()
120-
return -1
121-
elif unified == 'END WHILE':
122-
if self._block_stack and self._block_stack[-1] == 'WHILE':
123-
self._block_stack.pop()
124-
return -1
125-
elif unified == 'END LOOP':
126-
if self._block_stack and self._block_stack[-1] in ('LOOP', 'FOR', 'WHILE'):
127-
self._block_stack.pop()
128-
return -1
129-
elif unified == 'END CASE':
130-
if self._block_stack and self._block_stack[-1] == 'CASE':
131-
self._block_stack.pop()
132-
return -1
133-
elif unified == 'END':
134-
if self._block_stack:
135-
if self._block_stack[-1] in ('CASE', 'BEGIN'):
136-
self._block_stack.pop()
137-
return -1
138-
else:
139-
self._block_stack.pop()
140-
return -1
141-
else:
142-
return -1
143-
144-
return 0
148+
return self._handle_closing_keyword(unified)
145149

146150
def process(self, stream):
147151
"""Process the stream"""
@@ -173,12 +177,6 @@ def process(self, stream):
173177
# Issue809: Ignore semicolons inside BEGIN...END blocks, but handle
174178
# standalone BEGIN; as a transaction statement
175179
if ttype is T.Punctuation and value == ';':
176-
# If we just saw BEGIN; then this is a transaction BEGIN,
177-
# not a BEGIN...END block, so decrement depth
178-
if self._seen_begin:
179-
if self._block_stack and self._block_stack[-1] == 'BEGIN':
180-
self._block_stack.pop()
181-
self.level = max(0, self.level - 1)
182180
self._seen_begin = False
183181
# Split on semicolon if not inside a BEGIN...END block
184182
if self.level <= 0 and 'BEGIN' not in self._block_stack:

0 commit comments

Comments
 (0)