Repository navigation
Expand file tree
/
Copy pathstatement_splitter.py
More file actions
196 lines (171 loc) · 7.55 KB
/
Copy pathstatement_splitter.py
File metadata and controls
196 lines (171 loc) · 7.55 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
#
# Copyright (C) 2009-2020 the sqlparse authors and contributors
# <see AUTHORS file>
#
# This module is part of python-sqlparse and is released under
# the BSD License: https://opensource.org/licenses/BSD-3-Clause
from sqlparse import sql
from sqlparse import tokens as T
class StatementSplitter:
"""Filter that split stream at individual statements"""
def __init__(self):
self._reset()
def _reset(self):
"""Set the filter attributes to its default values"""
self._block_stack = []
self._parenthesis_level = 0
self._unconfirmed_start = None
self._is_create = False
self._seen_begin = False
self.consume_ws = False
self.tokens = []
self.level = 0
def _handle_nested_block(self, unified):
"""Check for nested loop or control structures inside a block"""
if unified == 'FOR':
self._unconfirmed_start = 'FOR'
return 0
if unified == 'WHILE':
self._unconfirmed_start = 'WHILE'
return 0
if unified in ('LOOP', 'DO'):
if self._unconfirmed_start in ('FOR', 'WHILE'):
self._block_stack.append(self._unconfirmed_start)
self._unconfirmed_start = None
return 1
if unified == 'LOOP':
self._block_stack.append('LOOP')
return 1
if unified in ('IF', 'CASE'):
self._block_stack.append(unified)
return 1
return None
def _handle_closing_keyword(self, unified):
"""Handle closing keywords for blocks"""
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:
self._block_stack.pop()
return -1
return 0
def _change_splitlevel(self, ttype, value):
"""Get the new split level (increase, decrease or remain equal)"""
# Semicolon resets unconfirmed loop starters
# and handles standalone BEGIN;
if ttype is T.Punctuation and value == ';':
self._unconfirmed_start = None
if self._seen_begin:
self._seen_begin = False
if self._block_stack and self._block_stack[-1] == 'BEGIN':
self._block_stack.pop()
return -1
return 0
# 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
unified = value.upper()
# 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
# 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._seen_begin = True
# 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
# Issue826: If we see a transaction keyword after BEGIN,
# it's a transaction statement, not a block.
if self._seen_begin and \
(ttype is T.Keyword or ttype is T.Name) and \
unified in ('TRANSACTION', 'WORK', 'TRAN',
'DISTRIBUTED', 'DEFERRED',
'IMMEDIATE', 'EXCLUSIVE'):
self._seen_begin = False
if self._block_stack and self._block_stack[-1] == 'BEGIN':
self._block_stack.pop()
return -1
return 0
# Inside a block, check for nested loop or control structures
if 'BEGIN' in self._block_stack:
res = self._handle_nested_block(unified)
if res is not None:
return res
# Handle closing keywords
return self._handle_closing_keyword(unified)
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:
yield sql.Statement(self.tokens)
# Reset filter and prepare to process next statement
self._reset()
# Change current split level (increase, decrease or remain equal)
self.level += self._change_splitlevel(ttype, value)
# Append the token to the current statement
self.tokens.append(sql.Token(ttype, value))
# Check if we get the end of a statement
# Issue762: Allow GO (or "GO 2") as statement splitter.
# When implementing a language toggle, it's not only to add
# keywords it's also to change some rules, like this splitting
# rule.
# Issue809: Ignore semicolons inside BEGIN...END blocks, but handle
# standalone BEGIN; as a transaction statement
if ttype is T.Punctuation and value == ';':
self._seen_begin = False
# Split on semicolon if not inside a BEGIN...END block
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
elif (ttype not in (T.Whitespace, T.Newline, T.Comment.Single,
T.Comment.Multiline)
and not (ttype is T.Keyword and value.upper() == 'BEGIN')):
# Reset _seen_begin if we see a non-whitespace, non-comment
# token but not for BEGIN itself (which just set the flag)
self._seen_begin = False
# Yield pending statement (if any)
if self.tokens and not all(t.is_whitespace for t in self.tokens):
yield sql.Statement(self.tokens)