1 |
2211
|
aaronmk
|
# SQL code generation
|
2 |
|
|
|
3 |
2748
|
aaronmk
|
import copy
|
4 |
3518
|
aaronmk
|
import itertools
|
5 |
2276
|
aaronmk
|
import operator
|
6 |
3158
|
aaronmk
|
from ordereddict import OrderedDict
|
7 |
2568
|
aaronmk
|
import re
|
8 |
2653
|
aaronmk
|
import UserDict
|
9 |
2953
|
aaronmk
|
import warnings
|
10 |
2276
|
aaronmk
|
|
11 |
2667
|
aaronmk
|
import dicts
|
12 |
2953
|
aaronmk
|
import exc
|
13 |
2701
|
aaronmk
|
import iters
|
14 |
|
|
import lists
|
15 |
2360
|
aaronmk
|
import objects
|
16 |
2222
|
aaronmk
|
import strings
|
17 |
2227
|
aaronmk
|
import util
|
18 |
2211
|
aaronmk
|
|
19 |
2587
|
aaronmk
|
##### Names
|
20 |
2499
|
aaronmk
|
|
21 |
2608
|
aaronmk
|
identifier_max_len = 63 # works for both PostgreSQL and MySQL
|
22 |
2587
|
aaronmk
|
|
23 |
2932
|
aaronmk
|
def concat(str_, suffix):
|
24 |
2609
|
aaronmk
|
'''Preserves version so that it won't be truncated off the string, leading
|
25 |
|
|
to collisions.'''
|
26 |
2613
|
aaronmk
|
# Preserve version
|
27 |
2995
|
aaronmk
|
match = re.match(r'^(.*?)((?:(?:#\d+)?\)?)*(?:\.\w+)?(?:::[\w ]+)*)$', str_)
|
28 |
2985
|
aaronmk
|
if match:
|
29 |
|
|
str_, old_suffix = match.groups()
|
30 |
|
|
suffix = old_suffix+suffix
|
31 |
2613
|
aaronmk
|
|
32 |
2932
|
aaronmk
|
return strings.concat(str_, suffix, identifier_max_len)
|
33 |
2587
|
aaronmk
|
|
34 |
2932
|
aaronmk
|
def truncate(str_): return concat(str_, '')
|
35 |
2842
|
aaronmk
|
|
36 |
2575
|
aaronmk
|
def is_safe_name(name):
|
37 |
2583
|
aaronmk
|
'''A name is safe *and unambiguous* if it:
|
38 |
|
|
* contains only *lowercase* word (\w) characters
|
39 |
|
|
* doesn't start with a digit
|
40 |
|
|
* contains "_", so that it's not a keyword
|
41 |
2984
|
aaronmk
|
'''
|
42 |
|
|
return re.match(r'^(?=.*_)(?!\d)[^\WA-Z]+$', name)
|
43 |
2568
|
aaronmk
|
|
44 |
2499
|
aaronmk
|
def esc_name(name, quote='"'):
|
45 |
|
|
return quote + name.replace(quote, quote+quote) + quote
|
46 |
|
|
# doubling an embedded quote escapes it in both PostgreSQL and MySQL
|
47 |
|
|
|
48 |
3320
|
aaronmk
|
def unesc_name(name, quote='"'):
|
49 |
|
|
removed_ref = [False]
|
50 |
|
|
name = strings.remove_prefix(quote, name, removed_ref)
|
51 |
|
|
if removed_ref[0]:
|
52 |
|
|
name = strings.remove_suffix(quote, name, removed_ref)
|
53 |
|
|
assert removed_ref[0]
|
54 |
|
|
name = name.replace(quote+quote, quote)
|
55 |
|
|
return name
|
56 |
|
|
|
57 |
2513
|
aaronmk
|
def clean_name(name): return name.replace('"', '').replace('`', '')
|
58 |
|
|
|
59 |
3041
|
aaronmk
|
def esc_comment(comment): return '/*'+comment.replace('*/', '* /')+'*/'
|
60 |
|
|
|
61 |
3182
|
aaronmk
|
def lstrip(str_):
|
62 |
|
|
'''Also removes comments.'''
|
63 |
|
|
if str_.startswith('/*'): comment, sep, str_ = str_.partition('*/')
|
64 |
|
|
return str_.lstrip()
|
65 |
|
|
|
66 |
2659
|
aaronmk
|
##### General SQL code objects
|
67 |
2219
|
aaronmk
|
|
68 |
2349
|
aaronmk
|
class MockDb:
|
69 |
2503
|
aaronmk
|
def esc_value(self, value): return strings.repr_no_u(value)
|
70 |
2349
|
aaronmk
|
|
71 |
2499
|
aaronmk
|
def esc_name(self, name): return esc_name(name)
|
72 |
2859
|
aaronmk
|
|
73 |
|
|
def col_info(self, col):
|
74 |
|
|
return TypedCol(col.name, '<type>', CustomCode('<default>'), True)
|
75 |
|
|
|
76 |
2349
|
aaronmk
|
mockDb = MockDb()
|
77 |
|
|
|
78 |
2514
|
aaronmk
|
class BasicObject(objects.BasicObject):
|
79 |
|
|
def __str__(self): return clean_name(strings.repr_no_u(self))
|
80 |
|
|
|
81 |
2659
|
aaronmk
|
##### Unparameterized code objects
|
82 |
|
|
|
83 |
2514
|
aaronmk
|
class Code(BasicObject):
|
84 |
3446
|
aaronmk
|
def __init__(self, lang='sql'):
|
85 |
|
|
self.lang = lang
|
86 |
3445
|
aaronmk
|
|
87 |
2658
|
aaronmk
|
def to_str(self, db): raise NotImplementedError()
|
88 |
2349
|
aaronmk
|
|
89 |
2514
|
aaronmk
|
def __repr__(self): return self.to_str(mockDb)
|
90 |
2211
|
aaronmk
|
|
91 |
2269
|
aaronmk
|
class CustomCode(Code):
|
92 |
3445
|
aaronmk
|
def __init__(self, str_):
|
93 |
|
|
Code.__init__(self)
|
94 |
|
|
|
95 |
|
|
self.str_ = str_
|
96 |
2256
|
aaronmk
|
|
97 |
|
|
def to_str(self, db): return self.str_
|
98 |
|
|
|
99 |
2815
|
aaronmk
|
def as_Code(value, db=None):
|
100 |
|
|
'''
|
101 |
|
|
@param db If set, runs db.std_code() on the value.
|
102 |
|
|
'''
|
103 |
3447
|
aaronmk
|
if isinstance(value, Code): return value
|
104 |
|
|
|
105 |
2815
|
aaronmk
|
if util.is_str(value):
|
106 |
|
|
if db != None: value = db.std_code(value)
|
107 |
|
|
return CustomCode(value)
|
108 |
2659
|
aaronmk
|
else: return Literal(value)
|
109 |
|
|
|
110 |
2540
|
aaronmk
|
class Expr(Code):
|
111 |
3445
|
aaronmk
|
def __init__(self, expr):
|
112 |
|
|
Code.__init__(self)
|
113 |
|
|
|
114 |
|
|
self.expr = expr
|
115 |
2540
|
aaronmk
|
|
116 |
|
|
def to_str(self, db): return '('+self.expr.to_str(db)+')'
|
117 |
|
|
|
118 |
3086
|
aaronmk
|
##### Names
|
119 |
|
|
|
120 |
|
|
class Name(Code):
|
121 |
3087
|
aaronmk
|
def __init__(self, name):
|
122 |
3445
|
aaronmk
|
Code.__init__(self)
|
123 |
|
|
|
124 |
3087
|
aaronmk
|
name = truncate(name)
|
125 |
|
|
|
126 |
|
|
self.name = name
|
127 |
3086
|
aaronmk
|
|
128 |
|
|
def to_str(self, db): return db.esc_name(self.name)
|
129 |
|
|
|
130 |
|
|
def as_Name(value):
|
131 |
|
|
if isinstance(value, Code): return value
|
132 |
|
|
else: return Name(value)
|
133 |
|
|
|
134 |
2335
|
aaronmk
|
##### Literal values
|
135 |
|
|
|
136 |
3519
|
aaronmk
|
#### Primitives
|
137 |
|
|
|
138 |
2216
|
aaronmk
|
class Literal(Code):
|
139 |
3445
|
aaronmk
|
def __init__(self, value):
|
140 |
|
|
Code.__init__(self)
|
141 |
|
|
|
142 |
|
|
self.value = value
|
143 |
2213
|
aaronmk
|
|
144 |
|
|
def to_str(self, db): return db.esc_value(self.value)
|
145 |
2211
|
aaronmk
|
|
146 |
2400
|
aaronmk
|
def as_Value(value):
|
147 |
|
|
if isinstance(value, Code): return value
|
148 |
|
|
else: return Literal(value)
|
149 |
|
|
|
150 |
3429
|
aaronmk
|
def is_literal(value): return isinstance(value, Literal)
|
151 |
2216
|
aaronmk
|
|
152 |
3429
|
aaronmk
|
def is_null(value): return is_literal(value) and value.value == None
|
153 |
|
|
|
154 |
3519
|
aaronmk
|
#### Composites
|
155 |
|
|
|
156 |
3521
|
aaronmk
|
class List(Code):
|
157 |
|
|
def __init__(self, values):
|
158 |
3519
|
aaronmk
|
Code.__init__(self)
|
159 |
|
|
|
160 |
|
|
self.values = values
|
161 |
|
|
|
162 |
3521
|
aaronmk
|
def to_str(self, db): return ', '.join((v.to_str(db) for v in self.values))
|
163 |
3519
|
aaronmk
|
|
164 |
3521
|
aaronmk
|
class Tuple(List):
|
165 |
|
|
def __init__(self, *values):
|
166 |
|
|
List.__init__(self, values)
|
167 |
|
|
|
168 |
|
|
def to_str(self, db): return '('+List.to_str(self, db)+')'
|
169 |
|
|
|
170 |
3520
|
aaronmk
|
class Row(Tuple):
|
171 |
|
|
def to_str(self, db): return 'ROW'+Tuple.to_str(self, db)
|
172 |
3519
|
aaronmk
|
|
173 |
3522
|
aaronmk
|
### Arrays
|
174 |
|
|
|
175 |
|
|
class Array(List):
|
176 |
|
|
def __init__(self, values):
|
177 |
|
|
values = map(remove_col_rename, values)
|
178 |
|
|
|
179 |
|
|
List.__init__(self, values)
|
180 |
|
|
|
181 |
|
|
def to_str(self, db): return 'ARRAY['+List.to_str(self, db)+']'
|
182 |
|
|
|
183 |
|
|
def to_Array(value):
|
184 |
|
|
if isinstance(value, Array): return value
|
185 |
|
|
return Array(lists.mk_seq(value))
|
186 |
|
|
|
187 |
2711
|
aaronmk
|
##### Derived elements
|
188 |
|
|
|
189 |
|
|
src_self = object() # tells Col that it is its own source column
|
190 |
|
|
|
191 |
|
|
class Derived(Code):
|
192 |
|
|
def __init__(self, srcs):
|
193 |
2712
|
aaronmk
|
'''An element which was derived from some other element(s).
|
194 |
2711
|
aaronmk
|
@param srcs See self.set_srcs()
|
195 |
|
|
'''
|
196 |
3445
|
aaronmk
|
Code.__init__(self)
|
197 |
|
|
|
198 |
2711
|
aaronmk
|
self.set_srcs(srcs)
|
199 |
|
|
|
200 |
2713
|
aaronmk
|
def set_srcs(self, srcs, overwrite=True):
|
201 |
2711
|
aaronmk
|
'''
|
202 |
|
|
@param srcs (self_type...)|src_self The element(s) this is derived from
|
203 |
|
|
'''
|
204 |
2713
|
aaronmk
|
if not overwrite and self.srcs != (): return # already set
|
205 |
|
|
|
206 |
2711
|
aaronmk
|
if srcs == src_self: srcs = (self,)
|
207 |
|
|
srcs = tuple(srcs) # make Col hashable
|
208 |
|
|
self.srcs = srcs
|
209 |
|
|
|
210 |
|
|
def _compare_on(self):
|
211 |
|
|
compare_on = self.__dict__.copy()
|
212 |
|
|
del compare_on['srcs'] # ignore
|
213 |
|
|
return compare_on
|
214 |
|
|
|
215 |
|
|
def cols_srcs(cols): return lists.uniqify(iters.flatten((v.srcs for v in cols)))
|
216 |
|
|
|
217 |
2335
|
aaronmk
|
##### Tables
|
218 |
|
|
|
219 |
2712
|
aaronmk
|
class Table(Derived):
|
220 |
2991
|
aaronmk
|
def __init__(self, name, schema=None, srcs=(), is_temp=False):
|
221 |
2211
|
aaronmk
|
'''
|
222 |
|
|
@param schema str|None (for no schema)
|
223 |
2712
|
aaronmk
|
@param srcs (Table...)|src_self See Derived.set_srcs()
|
224 |
2211
|
aaronmk
|
'''
|
225 |
2712
|
aaronmk
|
Derived.__init__(self, srcs)
|
226 |
|
|
|
227 |
3091
|
aaronmk
|
if util.is_str(name): name = truncate(name)
|
228 |
2843
|
aaronmk
|
|
229 |
2211
|
aaronmk
|
self.name = name
|
230 |
|
|
self.schema = schema
|
231 |
2991
|
aaronmk
|
self.is_temp = is_temp
|
232 |
3000
|
aaronmk
|
self.index_cols = {}
|
233 |
2211
|
aaronmk
|
|
234 |
2348
|
aaronmk
|
def to_str(self, db):
|
235 |
|
|
str_ = ''
|
236 |
3088
|
aaronmk
|
if self.schema != None: str_ += as_Name(self.schema).to_str(db)+'.'
|
237 |
|
|
str_ += as_Name(self.name).to_str(db)
|
238 |
2348
|
aaronmk
|
return str_
|
239 |
2336
|
aaronmk
|
|
240 |
|
|
def to_Table(self): return self
|
241 |
3000
|
aaronmk
|
|
242 |
|
|
def _compare_on(self):
|
243 |
|
|
compare_on = Derived._compare_on(self)
|
244 |
|
|
del compare_on['index_cols'] # ignore
|
245 |
|
|
return compare_on
|
246 |
2211
|
aaronmk
|
|
247 |
2835
|
aaronmk
|
def is_underlying_table(table):
|
248 |
|
|
return isinstance(table, Table) and table.to_Table() is table
|
249 |
2832
|
aaronmk
|
|
250 |
4114
|
aaronmk
|
class NoUnderlyingTableException(Exception):
|
251 |
|
|
def __init__(self, ref):
|
252 |
|
|
Exception.__init__(self, 'for: '+strings.as_tt(strings.urepr(ref)))
|
253 |
|
|
self.ref = ref
|
254 |
2902
|
aaronmk
|
|
255 |
|
|
def underlying_table(table):
|
256 |
|
|
table = remove_table_rename(table)
|
257 |
3532
|
aaronmk
|
if table != None and table.srcs:
|
258 |
|
|
table, = table.srcs # for derived tables or row vars
|
259 |
4114
|
aaronmk
|
if not is_underlying_table(table): raise NoUnderlyingTableException(table)
|
260 |
2902
|
aaronmk
|
return table
|
261 |
|
|
|
262 |
2776
|
aaronmk
|
def as_Table(table, schema=None):
|
263 |
2270
|
aaronmk
|
if table == None or isinstance(table, Code): return table
|
264 |
2776
|
aaronmk
|
else: return Table(table, schema)
|
265 |
2219
|
aaronmk
|
|
266 |
3101
|
aaronmk
|
def suffixed_table(table, suffix):
|
267 |
3128
|
aaronmk
|
table = copy.copy(table) # don't modify input!
|
268 |
|
|
table.name = concat(table.name, suffix)
|
269 |
|
|
return table
|
270 |
2707
|
aaronmk
|
|
271 |
2336
|
aaronmk
|
class NamedTable(Table):
|
272 |
|
|
def __init__(self, name, code, cols=None):
|
273 |
|
|
Table.__init__(self, name)
|
274 |
|
|
|
275 |
3016
|
aaronmk
|
code = as_Table(code)
|
276 |
2741
|
aaronmk
|
if not isinstance(code, (Table, FunctionCall, Expr)): code = Expr(code)
|
277 |
3020
|
aaronmk
|
if cols != None: cols = [to_name_only_col(c).to_Col() for c in cols]
|
278 |
2336
|
aaronmk
|
|
279 |
|
|
self.code = code
|
280 |
|
|
self.cols = cols
|
281 |
|
|
|
282 |
|
|
def to_str(self, db):
|
283 |
3026
|
aaronmk
|
str_ = self.code.to_str(db)
|
284 |
|
|
if str_.find('\n') >= 0: whitespace = '\n'
|
285 |
|
|
else: whitespace = ' '
|
286 |
|
|
str_ += whitespace+'AS '+Table.to_str(self, db)
|
287 |
2742
|
aaronmk
|
if self.cols != None:
|
288 |
|
|
str_ += ' ('+(', '.join((c.to_str(db) for c in self.cols)))+')'
|
289 |
2336
|
aaronmk
|
return str_
|
290 |
|
|
|
291 |
|
|
def to_Table(self): return Table(self.name)
|
292 |
|
|
|
293 |
2753
|
aaronmk
|
def remove_table_rename(table):
|
294 |
|
|
if isinstance(table, NamedTable): table = table.code
|
295 |
|
|
return table
|
296 |
|
|
|
297 |
2335
|
aaronmk
|
##### Columns
|
298 |
|
|
|
299 |
2711
|
aaronmk
|
class Col(Derived):
|
300 |
2701
|
aaronmk
|
def __init__(self, name, table=None, srcs=()):
|
301 |
2211
|
aaronmk
|
'''
|
302 |
|
|
@param table Table|None (for no table)
|
303 |
2711
|
aaronmk
|
@param srcs (Col...)|src_self See Derived.set_srcs()
|
304 |
2211
|
aaronmk
|
'''
|
305 |
2711
|
aaronmk
|
Derived.__init__(self, srcs)
|
306 |
|
|
|
307 |
3091
|
aaronmk
|
if util.is_str(name): name = truncate(name)
|
308 |
2241
|
aaronmk
|
if util.is_str(table): table = Table(table)
|
309 |
2211
|
aaronmk
|
assert table == None or isinstance(table, Table)
|
310 |
|
|
|
311 |
|
|
self.name = name
|
312 |
|
|
self.table = table
|
313 |
|
|
|
314 |
2989
|
aaronmk
|
def to_str(self, db, for_str=False):
|
315 |
3088
|
aaronmk
|
str_ = as_Name(self.name).to_str(db)
|
316 |
2989
|
aaronmk
|
if for_str: str_ = clean_name(str_)
|
317 |
2933
|
aaronmk
|
if self.table != None:
|
318 |
2989
|
aaronmk
|
table = self.table.to_Table()
|
319 |
3750
|
aaronmk
|
if for_str: str_ = concat(strings.ustr(table), '.'+str_)
|
320 |
2989
|
aaronmk
|
else: str_ = table.to_str(db)+'.'+str_
|
321 |
2211
|
aaronmk
|
return str_
|
322 |
2314
|
aaronmk
|
|
323 |
2989
|
aaronmk
|
def __str__(self): return self.to_str(mockDb, for_str=True)
|
324 |
2933
|
aaronmk
|
|
325 |
2314
|
aaronmk
|
def to_Col(self): return self
|
326 |
2211
|
aaronmk
|
|
327 |
3512
|
aaronmk
|
def is_col(col): return isinstance(col, Col)
|
328 |
2393
|
aaronmk
|
|
329 |
3512
|
aaronmk
|
def is_table_col(col): return is_col(col) and col.table != None
|
330 |
|
|
|
331 |
3000
|
aaronmk
|
def index_col(col):
|
332 |
|
|
if not is_table_col(col): return None
|
333 |
3104
|
aaronmk
|
|
334 |
|
|
table = col.table
|
335 |
|
|
try: name = table.index_cols[col.name]
|
336 |
|
|
except KeyError: return None
|
337 |
|
|
else: return Col(name, table, col.srcs)
|
338 |
2999
|
aaronmk
|
|
339 |
3024
|
aaronmk
|
def is_temp_col(col): return is_table_col(col) and col.table.is_temp
|
340 |
2996
|
aaronmk
|
|
341 |
2563
|
aaronmk
|
def as_Col(col, table=None, name=None):
|
342 |
|
|
'''
|
343 |
|
|
@param name If not None, any non-Col input will be renamed using NamedCol.
|
344 |
|
|
'''
|
345 |
|
|
if name != None:
|
346 |
|
|
col = as_Value(col)
|
347 |
|
|
if not isinstance(col, Col): col = NamedCol(name, col)
|
348 |
2333
|
aaronmk
|
|
349 |
|
|
if isinstance(col, Code): return col
|
350 |
3513
|
aaronmk
|
elif util.is_str(col): return Col(col, table)
|
351 |
|
|
else: return Literal(col)
|
352 |
2260
|
aaronmk
|
|
353 |
3093
|
aaronmk
|
def with_table(col, table):
|
354 |
3105
|
aaronmk
|
if isinstance(col, NamedCol): pass # doesn't take a table
|
355 |
|
|
elif isinstance(col, FunctionCall):
|
356 |
|
|
col = copy.deepcopy(col) # don't modify input!
|
357 |
|
|
col.args[0].table = table
|
358 |
3493
|
aaronmk
|
elif isinstance(col, Col):
|
359 |
3098
|
aaronmk
|
col = copy.copy(col) # don't modify input!
|
360 |
|
|
col.table = table
|
361 |
3093
|
aaronmk
|
return col
|
362 |
|
|
|
363 |
3100
|
aaronmk
|
def with_default_table(col, table):
|
364 |
2747
|
aaronmk
|
col = as_Col(col)
|
365 |
3100
|
aaronmk
|
if col.table == None: col = with_table(col, table)
|
366 |
2747
|
aaronmk
|
return col
|
367 |
|
|
|
368 |
2744
|
aaronmk
|
def set_cols_table(table, cols):
|
369 |
|
|
table = as_Table(table)
|
370 |
|
|
|
371 |
|
|
for i, col in enumerate(cols):
|
372 |
|
|
col = cols[i] = as_Col(col)
|
373 |
|
|
col.table = table
|
374 |
|
|
|
375 |
2401
|
aaronmk
|
def to_name_only_col(col, check_table=None):
|
376 |
|
|
col = as_Col(col)
|
377 |
3020
|
aaronmk
|
if not is_table_col(col): return col
|
378 |
2401
|
aaronmk
|
|
379 |
|
|
if check_table != None:
|
380 |
|
|
table = col.table
|
381 |
|
|
assert table == None or table == check_table
|
382 |
|
|
return Col(col.name)
|
383 |
|
|
|
384 |
2993
|
aaronmk
|
def suffixed_col(col, suffix):
|
385 |
|
|
return Col(concat(col.name, suffix), col.table, col.srcs)
|
386 |
|
|
|
387 |
3516
|
aaronmk
|
def has_srcs(col): return is_col(col) and col.srcs
|
388 |
|
|
|
389 |
3518
|
aaronmk
|
def cross_join_srcs(cols):
|
390 |
|
|
cols = filter(has_srcs, cols) # empty srcs will mess up the cross join
|
391 |
|
|
srcs = [[s.name for s in c.srcs] for c in cols]
|
392 |
|
|
return [Col(','.join(s)) for s in itertools.product(*srcs)]
|
393 |
3514
|
aaronmk
|
|
394 |
2323
|
aaronmk
|
class NamedCol(Col):
|
395 |
2229
|
aaronmk
|
def __init__(self, name, code):
|
396 |
2310
|
aaronmk
|
Col.__init__(self, name)
|
397 |
|
|
|
398 |
3016
|
aaronmk
|
code = as_Value(code)
|
399 |
2229
|
aaronmk
|
|
400 |
|
|
self.code = code
|
401 |
|
|
|
402 |
|
|
def to_str(self, db):
|
403 |
2310
|
aaronmk
|
return self.code.to_str(db)+' AS '+Col.to_str(self, db)
|
404 |
2314
|
aaronmk
|
|
405 |
|
|
def to_Col(self): return Col(self.name)
|
406 |
2229
|
aaronmk
|
|
407 |
2462
|
aaronmk
|
def remove_col_rename(col):
|
408 |
|
|
if isinstance(col, NamedCol): col = col.code
|
409 |
|
|
return col
|
410 |
|
|
|
411 |
2830
|
aaronmk
|
def underlying_col(col):
|
412 |
|
|
col = remove_col_rename(col)
|
413 |
4114
|
aaronmk
|
if not isinstance(col, Col): raise NoUnderlyingTableException(col)
|
414 |
2849
|
aaronmk
|
|
415 |
2902
|
aaronmk
|
return Col(col.name, underlying_table(col.table), col.srcs)
|
416 |
2830
|
aaronmk
|
|
417 |
2703
|
aaronmk
|
def wrap(wrap_func, value):
|
418 |
|
|
'''Wraps a value, propagating any column renaming to the returned value.'''
|
419 |
|
|
if isinstance(value, NamedCol):
|
420 |
|
|
return NamedCol(value.name, wrap_func(value.code))
|
421 |
|
|
else: return wrap_func(value)
|
422 |
|
|
|
423 |
2667
|
aaronmk
|
class ColDict(dicts.DictProxy):
|
424 |
3691
|
aaronmk
|
'''A dict that automatically makes inserted entries Col objects.
|
425 |
|
|
Anything that isn't a column is wrapped in a NamedCol with the key's column
|
426 |
|
|
name by `as_Col(value, name=key.name)`.
|
427 |
|
|
'''
|
428 |
2564
|
aaronmk
|
|
429 |
2645
|
aaronmk
|
def __init__(self, db, keys_table, dict_={}):
|
430 |
3158
|
aaronmk
|
dicts.DictProxy.__init__(self, OrderedDict())
|
431 |
2667
|
aaronmk
|
|
432 |
2645
|
aaronmk
|
keys_table = as_Table(keys_table)
|
433 |
|
|
|
434 |
2642
|
aaronmk
|
self.db = db
|
435 |
2641
|
aaronmk
|
self.table = keys_table
|
436 |
2653
|
aaronmk
|
self.update(dict_) # after setting vars because __setitem__() needs them
|
437 |
2641
|
aaronmk
|
|
438 |
2667
|
aaronmk
|
def copy(self): return ColDict(self.db, self.table, self.inner.copy())
|
439 |
2655
|
aaronmk
|
|
440 |
2667
|
aaronmk
|
def __getitem__(self, key):
|
441 |
|
|
return dicts.DictProxy.__getitem__(self, self._key(key))
|
442 |
2653
|
aaronmk
|
|
443 |
2564
|
aaronmk
|
def __setitem__(self, key, value):
|
444 |
2642
|
aaronmk
|
key = self._key(key)
|
445 |
3639
|
aaronmk
|
if value == None:
|
446 |
|
|
try: value = self.db.col_info(key).default
|
447 |
|
|
except NoUnderlyingTableException: pass # not a table column
|
448 |
2667
|
aaronmk
|
dicts.DictProxy.__setitem__(self, key, as_Col(value, name=key.name))
|
449 |
2564
|
aaronmk
|
|
450 |
2641
|
aaronmk
|
def _key(self, key): return as_Col(key, self.table)
|
451 |
2564
|
aaronmk
|
|
452 |
3519
|
aaronmk
|
##### Definitions
|
453 |
3491
|
aaronmk
|
|
454 |
3469
|
aaronmk
|
class TypedCol(Col):
|
455 |
|
|
def __init__(self, name, type_, default=None, nullable=True,
|
456 |
|
|
constraints=None):
|
457 |
|
|
assert default == None or isinstance(default, Code)
|
458 |
|
|
|
459 |
|
|
Col.__init__(self, name)
|
460 |
|
|
|
461 |
|
|
self.type = type_
|
462 |
|
|
self.default = default
|
463 |
|
|
self.nullable = nullable
|
464 |
|
|
self.constraints = constraints
|
465 |
|
|
|
466 |
|
|
def to_str(self, db):
|
467 |
3481
|
aaronmk
|
str_ = Col.to_str(self, db)+' '+as_Code(self.type).to_str(db)
|
468 |
3469
|
aaronmk
|
if not self.nullable: str_ += ' NOT NULL'
|
469 |
|
|
if self.default != None: str_ += ' DEFAULT '+self.default.to_str(db)
|
470 |
|
|
if self.constraints != None: str_ += ' '+self.constraints
|
471 |
|
|
return str_
|
472 |
|
|
|
473 |
|
|
def to_Col(self): return Col(self.name)
|
474 |
|
|
|
475 |
3488
|
aaronmk
|
class SetOf(Code):
|
476 |
|
|
def __init__(self, type_):
|
477 |
|
|
Code.__init__(self)
|
478 |
|
|
|
479 |
|
|
self.type = type_
|
480 |
|
|
|
481 |
|
|
def to_str(self, db):
|
482 |
|
|
return 'SETOF '+self.type.to_str(db)
|
483 |
|
|
|
484 |
3483
|
aaronmk
|
class RowType(Code):
|
485 |
|
|
def __init__(self, table):
|
486 |
|
|
Code.__init__(self)
|
487 |
|
|
|
488 |
|
|
self.table = table
|
489 |
|
|
|
490 |
|
|
def to_str(self, db):
|
491 |
|
|
return self.table.to_str(db)+'%ROWTYPE'
|
492 |
|
|
|
493 |
3485
|
aaronmk
|
class ColType(Code):
|
494 |
|
|
def __init__(self, col):
|
495 |
|
|
Code.__init__(self)
|
496 |
|
|
|
497 |
|
|
self.col = col
|
498 |
|
|
|
499 |
|
|
def to_str(self, db):
|
500 |
|
|
return self.col.to_str(db)+'%TYPE'
|
501 |
|
|
|
502 |
2524
|
aaronmk
|
##### Functions
|
503 |
|
|
|
504 |
2912
|
aaronmk
|
Function = Table
|
505 |
2911
|
aaronmk
|
as_Function = as_Table
|
506 |
|
|
|
507 |
2691
|
aaronmk
|
class InternalFunction(CustomCode): pass
|
508 |
|
|
|
509 |
3442
|
aaronmk
|
#### Calls
|
510 |
|
|
|
511 |
2941
|
aaronmk
|
class NamedArg(NamedCol):
|
512 |
|
|
def __init__(self, name, value):
|
513 |
|
|
NamedCol.__init__(self, name, value)
|
514 |
|
|
|
515 |
|
|
def to_str(self, db):
|
516 |
|
|
return Col.to_str(self, db)+' := '+self.code.to_str(db)
|
517 |
|
|
|
518 |
2524
|
aaronmk
|
class FunctionCall(Code):
|
519 |
2941
|
aaronmk
|
def __init__(self, function, *args, **kw_args):
|
520 |
2524
|
aaronmk
|
'''
|
521 |
2690
|
aaronmk
|
@param args [Code|literal-value...] The function's arguments
|
522 |
2524
|
aaronmk
|
'''
|
523 |
3445
|
aaronmk
|
Code.__init__(self)
|
524 |
|
|
|
525 |
3016
|
aaronmk
|
function = as_Function(function)
|
526 |
2941
|
aaronmk
|
def filter_(arg): return remove_col_rename(as_Value(arg))
|
527 |
|
|
args = map(filter_, args)
|
528 |
|
|
args += [NamedArg(k, filter_(v)) for k, v in kw_args.iteritems()]
|
529 |
2524
|
aaronmk
|
|
530 |
|
|
self.function = function
|
531 |
|
|
self.args = args
|
532 |
|
|
|
533 |
|
|
def to_str(self, db):
|
534 |
|
|
args_str = ', '.join((v.to_str(db) for v in self.args))
|
535 |
|
|
return self.function.to_str(db)+'('+args_str+')'
|
536 |
|
|
|
537 |
2533
|
aaronmk
|
def wrap_in_func(function, value):
|
538 |
|
|
'''Wraps a value inside a function call.
|
539 |
|
|
Propagates any column renaming to the returned value.
|
540 |
|
|
'''
|
541 |
2703
|
aaronmk
|
return wrap(lambda v: FunctionCall(function, v), value)
|
542 |
2533
|
aaronmk
|
|
543 |
2561
|
aaronmk
|
def unwrap_func_call(func_call, check_name=None):
|
544 |
|
|
'''Unwraps any function call to its first argument.
|
545 |
|
|
Also removes any column renaming.
|
546 |
|
|
'''
|
547 |
|
|
func_call = remove_col_rename(func_call)
|
548 |
|
|
if not isinstance(func_call, FunctionCall): return func_call
|
549 |
|
|
|
550 |
|
|
if check_name != None:
|
551 |
|
|
name = func_call.function.name
|
552 |
|
|
assert name == None or name == check_name
|
553 |
|
|
return func_call.args[0]
|
554 |
|
|
|
555 |
3442
|
aaronmk
|
#### Definitions
|
556 |
|
|
|
557 |
|
|
class FunctionDef(Code):
|
558 |
3471
|
aaronmk
|
def __init__(self, function, return_type, body, params=[], modifiers=None):
|
559 |
3445
|
aaronmk
|
Code.__init__(self)
|
560 |
|
|
|
561 |
3487
|
aaronmk
|
return_type = as_Code(return_type)
|
562 |
3444
|
aaronmk
|
body = as_Code(body)
|
563 |
|
|
|
564 |
3442
|
aaronmk
|
self.function = function
|
565 |
|
|
self.return_type = return_type
|
566 |
|
|
self.body = body
|
567 |
3471
|
aaronmk
|
self.params = params
|
568 |
3456
|
aaronmk
|
self.modifiers = modifiers
|
569 |
3442
|
aaronmk
|
|
570 |
|
|
def to_str(self, db):
|
571 |
3487
|
aaronmk
|
params_str = (', '.join((p.to_str(db) for p in self.params)))
|
572 |
3442
|
aaronmk
|
str_ = '''\
|
573 |
3471
|
aaronmk
|
CREATE FUNCTION '''+self.function.to_str(db)+'''('''+params_str+''')
|
574 |
3487
|
aaronmk
|
RETURNS '''+self.return_type.to_str(db)+'''
|
575 |
3448
|
aaronmk
|
LANGUAGE '''+self.body.lang+'''
|
576 |
3456
|
aaronmk
|
'''
|
577 |
|
|
if self.modifiers != None: str_ += self.modifiers+'\n'
|
578 |
|
|
str_ += '''\
|
579 |
3442
|
aaronmk
|
AS $$
|
580 |
3444
|
aaronmk
|
'''+self.body.to_str(db)+'''
|
581 |
3442
|
aaronmk
|
$$;
|
582 |
|
|
'''
|
583 |
|
|
return str_
|
584 |
|
|
|
585 |
3469
|
aaronmk
|
class FunctionParam(TypedCol):
|
586 |
|
|
def __init__(self, name, type_, default=None, out=False):
|
587 |
|
|
TypedCol.__init__(self, name, type_, default)
|
588 |
|
|
|
589 |
|
|
self.out = out
|
590 |
|
|
|
591 |
|
|
def to_str(self, db):
|
592 |
|
|
str_ = TypedCol.to_str(self, db)
|
593 |
|
|
if self.out: str_ = 'OUT '+str_
|
594 |
|
|
return str_
|
595 |
|
|
|
596 |
|
|
def to_Col(self): return Col(self.name)
|
597 |
|
|
|
598 |
3454
|
aaronmk
|
### PL/pgSQL
|
599 |
|
|
|
600 |
3496
|
aaronmk
|
class ReturnQuery(Code):
|
601 |
|
|
def __init__(self, query):
|
602 |
|
|
Code.__init__(self)
|
603 |
|
|
|
604 |
|
|
query = as_Code(query)
|
605 |
|
|
|
606 |
|
|
self.query = query
|
607 |
|
|
|
608 |
|
|
def to_str(self, db):
|
609 |
|
|
return 'RETURN QUERY\n'+strings.indent(self.query.to_str(db))+';\n'
|
610 |
|
|
|
611 |
3515
|
aaronmk
|
## Exceptions
|
612 |
|
|
|
613 |
3509
|
aaronmk
|
class BaseExcHandler(BasicObject):
|
614 |
|
|
def to_str(self, db, body): raise NotImplementedError()
|
615 |
3510
|
aaronmk
|
|
616 |
|
|
def __repr__(self): return self.to_str(mockDb, '<body>')
|
617 |
3509
|
aaronmk
|
|
618 |
3549
|
aaronmk
|
suppress_exc = 'NULL;\n';
|
619 |
|
|
|
620 |
3554
|
aaronmk
|
reraise_exc = 'RAISE USING ERRCODE = SQLSTATE, MESSAGE = SQLERRM;\n';
|
621 |
|
|
|
622 |
3509
|
aaronmk
|
class ExcHandler(BaseExcHandler):
|
623 |
3454
|
aaronmk
|
def __init__(self, exc, handler=None):
|
624 |
|
|
if handler != None: handler = as_Code(handler)
|
625 |
|
|
|
626 |
|
|
self.exc = exc
|
627 |
|
|
self.handler = handler
|
628 |
|
|
|
629 |
|
|
def to_str(self, db, body):
|
630 |
|
|
body = as_Code(body)
|
631 |
|
|
|
632 |
3467
|
aaronmk
|
if self.handler != None:
|
633 |
|
|
handler_str = '\n'+strings.indent(self.handler.to_str(db), 2)
|
634 |
3549
|
aaronmk
|
else: handler_str = ' '+suppress_exc
|
635 |
3454
|
aaronmk
|
|
636 |
|
|
str_ = '''\
|
637 |
|
|
BEGIN
|
638 |
3467
|
aaronmk
|
'''+strings.indent(body.to_str(db))+'''\
|
639 |
3454
|
aaronmk
|
EXCEPTION
|
640 |
3463
|
aaronmk
|
WHEN '''+self.exc+''' THEN'''+handler_str+'''\
|
641 |
3454
|
aaronmk
|
END;\
|
642 |
|
|
'''
|
643 |
|
|
return str_
|
644 |
|
|
|
645 |
3515
|
aaronmk
|
class NestedExcHandler(BaseExcHandler):
|
646 |
|
|
def __init__(self, *handlers):
|
647 |
|
|
'''
|
648 |
|
|
@param handlers Sorted from outermost to innermost
|
649 |
|
|
'''
|
650 |
|
|
self.handlers = handlers
|
651 |
|
|
|
652 |
|
|
def to_str(self, db, body):
|
653 |
|
|
for handler in reversed(self.handlers): body = handler.to_str(db, body)
|
654 |
|
|
return body
|
655 |
|
|
|
656 |
3503
|
aaronmk
|
class ExcToWarning(Code):
|
657 |
|
|
def __init__(self, return_):
|
658 |
|
|
'''
|
659 |
|
|
@param return_ Statement to return a default value in case of error
|
660 |
|
|
'''
|
661 |
|
|
Code.__init__(self)
|
662 |
|
|
|
663 |
|
|
return_ = as_Code(return_)
|
664 |
|
|
|
665 |
|
|
self.return_ = return_
|
666 |
|
|
|
667 |
|
|
def to_str(self, db):
|
668 |
|
|
return "RAISE WARNING '%', SQLERRM;\n"+self.return_.to_str(db)
|
669 |
|
|
|
670 |
3454
|
aaronmk
|
unique_violation_handler = ExcHandler('unique_violation')
|
671 |
|
|
|
672 |
3555
|
aaronmk
|
# Note doubled "\"s because inside Python string
|
673 |
3468
|
aaronmk
|
plpythonu_error_handler = ExcHandler('internal_error', '''\
|
674 |
3591
|
aaronmk
|
-- Handle PL/Python exceptions
|
675 |
3555
|
aaronmk
|
DECLARE
|
676 |
3595
|
aaronmk
|
matches text[] := regexp_matches(SQLERRM,
|
677 |
|
|
E'^(?:PL/Python: )?(\\\\w+): (.*)$'); -- .* also matches \\n
|
678 |
3555
|
aaronmk
|
exc_name text := matches[1];
|
679 |
|
|
msg text := matches[2];
|
680 |
|
|
BEGIN
|
681 |
3596
|
aaronmk
|
/* Re-raise PL/Python exceptions with the PL/Python prefix removed.
|
682 |
|
|
This allows the exception to be parsed like a native exception.
|
683 |
|
|
Always raise as data_exception so it goes in the errors table. */
|
684 |
3598
|
aaronmk
|
IF exc_name IS NOT NULL THEN
|
685 |
4026
|
aaronmk
|
RAISE data_exception USING MESSAGE = msg;
|
686 |
3595
|
aaronmk
|
-- Re-raise non-PL/Python exceptions
|
687 |
|
|
ELSE
|
688 |
|
|
'''+reraise_exc+'''\
|
689 |
3555
|
aaronmk
|
END IF;
|
690 |
|
|
END;
|
691 |
3468
|
aaronmk
|
''')
|
692 |
|
|
|
693 |
3505
|
aaronmk
|
def data_exception_handler(handler):
|
694 |
|
|
return ExcHandler('data_exception', handler)
|
695 |
|
|
|
696 |
3527
|
aaronmk
|
row_var = Table('row')
|
697 |
|
|
|
698 |
3449
|
aaronmk
|
class RowExcIgnore(Code):
|
699 |
|
|
def __init__(self, row_type, select_query, with_row, cols=None,
|
700 |
3527
|
aaronmk
|
exc_handler=unique_violation_handler, row_var=row_var):
|
701 |
3529
|
aaronmk
|
'''
|
702 |
|
|
@param row_type Ignored if a custom row_var is used.
|
703 |
|
|
@pre If a custom row_var is used, it must already be defined.
|
704 |
|
|
'''
|
705 |
3449
|
aaronmk
|
Code.__init__(self, lang='plpgsql')
|
706 |
|
|
|
707 |
3482
|
aaronmk
|
row_type = as_Code(row_type)
|
708 |
3449
|
aaronmk
|
select_query = as_Code(select_query)
|
709 |
|
|
with_row = as_Code(with_row)
|
710 |
3452
|
aaronmk
|
row_var = as_Table(row_var)
|
711 |
3449
|
aaronmk
|
|
712 |
|
|
self.row_type = row_type
|
713 |
|
|
self.select_query = select_query
|
714 |
|
|
self.with_row = with_row
|
715 |
|
|
self.cols = cols
|
716 |
3455
|
aaronmk
|
self.exc_handler = exc_handler
|
717 |
3452
|
aaronmk
|
self.row_var = row_var
|
718 |
3449
|
aaronmk
|
|
719 |
|
|
def to_str(self, db):
|
720 |
3452
|
aaronmk
|
if self.cols == None: row_vars = [self.row_var]
|
721 |
|
|
else: row_vars = [Col(c.name, self.row_var) for c in self.cols]
|
722 |
3449
|
aaronmk
|
|
723 |
3526
|
aaronmk
|
# Need an EXCEPTION block for each individual row because "When an error
|
724 |
|
|
# is caught by an EXCEPTION clause, [...] all changes to persistent
|
725 |
|
|
# database state within the block are rolled back."
|
726 |
|
|
# This is unfortunate because "A block containing an EXCEPTION clause is
|
727 |
|
|
# significantly more expensive to enter and exit than a block without
|
728 |
|
|
# one."
|
729 |
|
|
# (http://www.postgresql.org/docs/8.3/static/\
|
730 |
|
|
# plpgsql-control-structures.html#PLPGSQL-ERROR-TRAPPING)
|
731 |
3449
|
aaronmk
|
str_ = '''\
|
732 |
3529
|
aaronmk
|
FOR '''+(', '.join((v.to_str(db) for v in row_vars)))+''' IN
|
733 |
|
|
'''+strings.indent(self.select_query.to_str(db))+'''\
|
734 |
|
|
LOOP
|
735 |
|
|
'''+strings.indent(self.exc_handler.to_str(db, self.with_row))+'''\
|
736 |
|
|
END LOOP;
|
737 |
|
|
'''
|
738 |
3533
|
aaronmk
|
if self.row_var == row_var:
|
739 |
3529
|
aaronmk
|
str_ = '''\
|
740 |
3449
|
aaronmk
|
DECLARE
|
741 |
3482
|
aaronmk
|
'''+self.row_var.to_str(db)+''' '''+self.row_type.to_str(db)+''';
|
742 |
3449
|
aaronmk
|
BEGIN
|
743 |
3529
|
aaronmk
|
'''+strings.indent(str_)+'''\
|
744 |
|
|
END;
|
745 |
3449
|
aaronmk
|
'''
|
746 |
|
|
return str_
|
747 |
|
|
|
748 |
2986
|
aaronmk
|
##### Casts
|
749 |
|
|
|
750 |
|
|
class Cast(FunctionCall):
|
751 |
|
|
def __init__(self, type_, value):
|
752 |
3539
|
aaronmk
|
type_ = as_Code(type_)
|
753 |
2986
|
aaronmk
|
value = as_Value(value)
|
754 |
|
|
|
755 |
|
|
self.type_ = type_
|
756 |
|
|
self.value = value
|
757 |
|
|
|
758 |
|
|
def to_str(self, db):
|
759 |
3539
|
aaronmk
|
return 'CAST('+self.value.to_str(db)+' AS '+self.type_.to_str(db)+')'
|
760 |
2986
|
aaronmk
|
|
761 |
3354
|
aaronmk
|
def cast_literal(value):
|
762 |
3429
|
aaronmk
|
if not is_literal(value): return value
|
763 |
3354
|
aaronmk
|
|
764 |
|
|
if util.is_str(value.value): value = Cast('text', value)
|
765 |
|
|
return value
|
766 |
|
|
|
767 |
2335
|
aaronmk
|
##### Conditions
|
768 |
2259
|
aaronmk
|
|
769 |
3350
|
aaronmk
|
class NotCond(Code):
|
770 |
|
|
def __init__(self, cond):
|
771 |
3445
|
aaronmk
|
Code.__init__(self)
|
772 |
|
|
|
773 |
3350
|
aaronmk
|
self.cond = cond
|
774 |
|
|
|
775 |
|
|
def to_str(self, db): return 'NOT '+self.cond.to_str(db)
|
776 |
|
|
|
777 |
2398
|
aaronmk
|
class ColValueCond(Code):
|
778 |
|
|
def __init__(self, col, value):
|
779 |
3445
|
aaronmk
|
Code.__init__(self)
|
780 |
|
|
|
781 |
2398
|
aaronmk
|
value = as_ValueCond(value)
|
782 |
|
|
|
783 |
|
|
self.col = col
|
784 |
|
|
self.value = value
|
785 |
|
|
|
786 |
|
|
def to_str(self, db): return self.value.to_str(db, self.col)
|
787 |
|
|
|
788 |
2577
|
aaronmk
|
def combine_conds(conds, keyword=None):
|
789 |
|
|
'''
|
790 |
|
|
@param keyword The keyword to add before the conditions, if any
|
791 |
|
|
'''
|
792 |
|
|
str_ = ''
|
793 |
|
|
if keyword != None:
|
794 |
|
|
if conds == []: whitespace = ''
|
795 |
|
|
elif len(conds) == 1: whitespace = ' '
|
796 |
|
|
else: whitespace = '\n'
|
797 |
|
|
str_ += keyword+whitespace
|
798 |
|
|
|
799 |
|
|
str_ += '\nAND '.join(conds)
|
800 |
|
|
return str_
|
801 |
|
|
|
802 |
2398
|
aaronmk
|
##### Condition column comparisons
|
803 |
|
|
|
804 |
2514
|
aaronmk
|
class ValueCond(BasicObject):
|
805 |
2213
|
aaronmk
|
def __init__(self, value):
|
806 |
2858
|
aaronmk
|
value = remove_col_rename(as_Value(value))
|
807 |
2213
|
aaronmk
|
|
808 |
|
|
self.value = value
|
809 |
2214
|
aaronmk
|
|
810 |
2216
|
aaronmk
|
def to_str(self, db, left_value):
|
811 |
2214
|
aaronmk
|
'''
|
812 |
2216
|
aaronmk
|
@param left_value The Code object that the condition is being applied on
|
813 |
2214
|
aaronmk
|
'''
|
814 |
|
|
raise NotImplemented()
|
815 |
2228
|
aaronmk
|
|
816 |
2514
|
aaronmk
|
def __repr__(self): return self.to_str(mockDb, '<left_value>')
|
817 |
2211
|
aaronmk
|
|
818 |
|
|
class CompareCond(ValueCond):
|
819 |
|
|
def __init__(self, value, operator='='):
|
820 |
2222
|
aaronmk
|
'''
|
821 |
|
|
@param operator By default, compares NULL values literally. Use '~=' or
|
822 |
|
|
'~!=' to pass NULLs through.
|
823 |
|
|
'''
|
824 |
2211
|
aaronmk
|
ValueCond.__init__(self, value)
|
825 |
|
|
self.operator = operator
|
826 |
|
|
|
827 |
2216
|
aaronmk
|
def to_str(self, db, left_value):
|
828 |
2858
|
aaronmk
|
left_value = remove_col_rename(as_Col(left_value))
|
829 |
2216
|
aaronmk
|
|
830 |
2222
|
aaronmk
|
right_value = self.value
|
831 |
|
|
|
832 |
|
|
# Parse operator
|
833 |
2216
|
aaronmk
|
operator = self.operator
|
834 |
2222
|
aaronmk
|
passthru_null_ref = [False]
|
835 |
|
|
operator = strings.remove_prefix('~', operator, passthru_null_ref)
|
836 |
|
|
neg_ref = [False]
|
837 |
|
|
operator = strings.remove_prefix('!', operator, neg_ref)
|
838 |
2844
|
aaronmk
|
equals = operator.endswith('=') # also includes <=, >=
|
839 |
2222
|
aaronmk
|
|
840 |
2825
|
aaronmk
|
# Handle nullable columns
|
841 |
|
|
check_null = False
|
842 |
2844
|
aaronmk
|
if not passthru_null_ref[0]: # NULLs compare equal
|
843 |
2857
|
aaronmk
|
try: left_value = ensure_not_null(db, left_value)
|
844 |
2844
|
aaronmk
|
except ensure_not_null_excs: # fall back to alternate method
|
845 |
|
|
check_null = equals and isinstance(right_value, Col)
|
846 |
2837
|
aaronmk
|
else:
|
847 |
2857
|
aaronmk
|
if isinstance(left_value, EnsureNotNull):
|
848 |
|
|
right_value = ensure_not_null(db, right_value,
|
849 |
|
|
left_value.type) # apply same function to both sides
|
850 |
2825
|
aaronmk
|
|
851 |
2844
|
aaronmk
|
if equals and is_null(right_value): operator = 'IS'
|
852 |
|
|
|
853 |
2825
|
aaronmk
|
left = left_value.to_str(db)
|
854 |
|
|
right = right_value.to_str(db)
|
855 |
|
|
|
856 |
2222
|
aaronmk
|
# Create str
|
857 |
|
|
str_ = left+' '+operator+' '+right
|
858 |
2825
|
aaronmk
|
if check_null:
|
859 |
2578
|
aaronmk
|
str_ = '('+str_+' OR ('+left+' IS NULL AND '+right+' IS NULL))'
|
860 |
|
|
if neg_ref[0]: str_ = 'NOT '+str_
|
861 |
2222
|
aaronmk
|
return str_
|
862 |
2216
|
aaronmk
|
|
863 |
2260
|
aaronmk
|
# Tells as_ValueCond() to assume a non-ValueCond is a literal value
|
864 |
|
|
assume_literal = object()
|
865 |
|
|
|
866 |
|
|
def as_ValueCond(value, default_table=assume_literal):
|
867 |
|
|
if not isinstance(value, ValueCond):
|
868 |
|
|
if default_table is not assume_literal:
|
869 |
2748
|
aaronmk
|
value = with_default_table(value, default_table)
|
870 |
2260
|
aaronmk
|
return CompareCond(value)
|
871 |
2216
|
aaronmk
|
else: return value
|
872 |
2219
|
aaronmk
|
|
873 |
2335
|
aaronmk
|
##### Joins
|
874 |
|
|
|
875 |
2352
|
aaronmk
|
join_same = object() # tells Join the left and right columns have the same name
|
876 |
2260
|
aaronmk
|
|
877 |
2353
|
aaronmk
|
# Tells Join the left and right columns have the same name and are never NULL
|
878 |
|
|
join_same_not_null = object()
|
879 |
|
|
|
880 |
2260
|
aaronmk
|
filter_out = object() # tells Join to filter out rows that match the join
|
881 |
|
|
|
882 |
2514
|
aaronmk
|
class Join(BasicObject):
|
883 |
2746
|
aaronmk
|
def __init__(self, table, mapping={}, type_=None):
|
884 |
2260
|
aaronmk
|
'''
|
885 |
|
|
@param mapping dict(right_table_col=left_table_col, ...)
|
886 |
2352
|
aaronmk
|
* if left_table_col is join_same: left_table_col = right_table_col
|
887 |
2353
|
aaronmk
|
* Note that right_table_col must be a string
|
888 |
|
|
* if left_table_col is join_same_not_null:
|
889 |
|
|
left_table_col = right_table_col and both have NOT NULL constraint
|
890 |
|
|
* Note that right_table_col must be a string
|
891 |
2260
|
aaronmk
|
@param type_ None (for plain join)|str (e.g. 'LEFT')|filter_out
|
892 |
|
|
* filter_out: equivalent to 'LEFT' with the query filtered by
|
893 |
|
|
`table_pkey IS NULL` (indicating no match)
|
894 |
|
|
'''
|
895 |
|
|
if util.is_str(table): table = Table(table)
|
896 |
|
|
assert type_ == None or util.is_str(type_) or type_ is filter_out
|
897 |
|
|
|
898 |
|
|
self.table = table
|
899 |
|
|
self.mapping = mapping
|
900 |
|
|
self.type_ = type_
|
901 |
|
|
|
902 |
2749
|
aaronmk
|
def to_str(self, db, left_table_):
|
903 |
2260
|
aaronmk
|
def join(entry):
|
904 |
|
|
'''Parses non-USING joins'''
|
905 |
|
|
right_table_col, left_table_col = entry
|
906 |
|
|
|
907 |
2353
|
aaronmk
|
# Switch order (right_table_col is on the left in the comparison)
|
908 |
|
|
left = right_table_col
|
909 |
|
|
right = left_table_col
|
910 |
2749
|
aaronmk
|
left_table = self.table
|
911 |
|
|
right_table = left_table_
|
912 |
2353
|
aaronmk
|
|
913 |
2747
|
aaronmk
|
# Parse left side
|
914 |
2748
|
aaronmk
|
left = with_default_table(left, left_table)
|
915 |
2747
|
aaronmk
|
|
916 |
2260
|
aaronmk
|
# Parse special values
|
917 |
2747
|
aaronmk
|
left_on_right = Col(left.name, right_table)
|
918 |
|
|
if right is join_same: right = left_on_right
|
919 |
2353
|
aaronmk
|
elif right is join_same_not_null:
|
920 |
2747
|
aaronmk
|
right = CompareCond(left_on_right, '~=')
|
921 |
2260
|
aaronmk
|
|
922 |
2747
|
aaronmk
|
# Parse right side
|
923 |
2353
|
aaronmk
|
right = as_ValueCond(right, right_table)
|
924 |
2747
|
aaronmk
|
|
925 |
|
|
return right.to_str(db, left)
|
926 |
2260
|
aaronmk
|
|
927 |
2265
|
aaronmk
|
# Create join condition
|
928 |
|
|
type_ = self.type_
|
929 |
2276
|
aaronmk
|
joins = self.mapping
|
930 |
2746
|
aaronmk
|
if joins == {}: join_cond = None
|
931 |
|
|
elif type_ is not filter_out and reduce(operator.and_,
|
932 |
2460
|
aaronmk
|
(v is join_same_not_null for v in joins.itervalues())):
|
933 |
2260
|
aaronmk
|
# all cols w/ USING, so can use simpler USING syntax
|
934 |
2747
|
aaronmk
|
cols = map(to_name_only_col, joins.iterkeys())
|
935 |
|
|
join_cond = 'USING ('+(', '.join((c.to_str(db) for c in cols)))+')'
|
936 |
2757
|
aaronmk
|
else: join_cond = combine_conds(map(join, joins.iteritems()), 'ON')
|
937 |
2260
|
aaronmk
|
|
938 |
2757
|
aaronmk
|
if isinstance(self.table, NamedTable): whitespace = '\n'
|
939 |
|
|
else: whitespace = ' '
|
940 |
|
|
|
941 |
2260
|
aaronmk
|
# Create join
|
942 |
|
|
if type_ is filter_out: type_ = 'LEFT'
|
943 |
2266
|
aaronmk
|
str_ = ''
|
944 |
|
|
if type_ != None: str_ += type_+' '
|
945 |
2757
|
aaronmk
|
str_ += 'JOIN'+whitespace+self.table.to_str(db)
|
946 |
|
|
if join_cond != None: str_ += whitespace+join_cond
|
947 |
2266
|
aaronmk
|
return str_
|
948 |
2349
|
aaronmk
|
|
949 |
2514
|
aaronmk
|
def __repr__(self): return self.to_str(mockDb, '<left_table>')
|
950 |
2424
|
aaronmk
|
|
951 |
|
|
##### Value exprs
|
952 |
|
|
|
953 |
3089
|
aaronmk
|
all_cols = CustomCode('*')
|
954 |
|
|
|
955 |
2737
|
aaronmk
|
default = CustomCode('DEFAULT')
|
956 |
|
|
|
957 |
3090
|
aaronmk
|
row_count = FunctionCall(InternalFunction('COUNT'), all_cols)
|
958 |
2674
|
aaronmk
|
|
959 |
3061
|
aaronmk
|
class Coalesce(FunctionCall):
|
960 |
|
|
def __init__(self, *args):
|
961 |
|
|
FunctionCall.__init__(self, InternalFunction('COALESCE'), *args)
|
962 |
3060
|
aaronmk
|
|
963 |
3062
|
aaronmk
|
class Nullif(FunctionCall):
|
964 |
|
|
def __init__(self, *args):
|
965 |
|
|
FunctionCall.__init__(self, InternalFunction('NULLIF'), *args)
|
966 |
|
|
|
967 |
3706
|
aaronmk
|
null_as_str = Cast('text', 'NULL')
|
968 |
|
|
|
969 |
|
|
def to_text(value): return Coalesce(Cast('text', value), null_as_str)
|
970 |
|
|
|
971 |
2850
|
aaronmk
|
# See <http://www.postgresql.org/docs/8.3/static/datatype-numeric.html>
|
972 |
2958
|
aaronmk
|
null_sentinels = {
|
973 |
|
|
'character varying': r'\N',
|
974 |
|
|
'double precision': 'NaN',
|
975 |
|
|
'integer': 2147483647,
|
976 |
|
|
'text': r'\N',
|
977 |
|
|
'timestamp with time zone': 'infinity'
|
978 |
|
|
}
|
979 |
2692
|
aaronmk
|
|
980 |
3061
|
aaronmk
|
class EnsureNotNull(Coalesce):
|
981 |
2850
|
aaronmk
|
def __init__(self, value, type_):
|
982 |
3061
|
aaronmk
|
Coalesce.__init__(self, as_Col(value),
|
983 |
2988
|
aaronmk
|
Cast(type_, null_sentinels[type_]))
|
984 |
2850
|
aaronmk
|
|
985 |
|
|
self.type = type_
|
986 |
3001
|
aaronmk
|
|
987 |
|
|
def to_str(self, db):
|
988 |
|
|
col = self.args[0]
|
989 |
|
|
index_col_ = index_col(col)
|
990 |
|
|
if index_col_ != None: return index_col_.to_str(db)
|
991 |
3061
|
aaronmk
|
return Coalesce.to_str(self, db)
|
992 |
2850
|
aaronmk
|
|
993 |
3523
|
aaronmk
|
#### Arrays
|
994 |
|
|
|
995 |
3535
|
aaronmk
|
class ArrayMerge(FunctionCall):
|
996 |
3523
|
aaronmk
|
def __init__(self, sep, array):
|
997 |
|
|
array = to_Array(array)
|
998 |
|
|
FunctionCall.__init__(self, InternalFunction('array_to_string'), array,
|
999 |
|
|
sep)
|
1000 |
|
|
|
1001 |
3537
|
aaronmk
|
def merge_not_null(db, sep, values):
|
1002 |
3707
|
aaronmk
|
return ArrayMerge(sep, map(to_text, values))
|
1003 |
3537
|
aaronmk
|
|
1004 |
2737
|
aaronmk
|
##### Table exprs
|
1005 |
|
|
|
1006 |
|
|
class Values(Code):
|
1007 |
|
|
def __init__(self, values):
|
1008 |
2739
|
aaronmk
|
'''
|
1009 |
|
|
@param values [...]|[[...], ...] Can be one or multiple rows.
|
1010 |
|
|
'''
|
1011 |
3445
|
aaronmk
|
Code.__init__(self)
|
1012 |
|
|
|
1013 |
2739
|
aaronmk
|
rows = values
|
1014 |
|
|
if len(values) >= 1 and not lists.is_seq(values[0]): # only one row
|
1015 |
|
|
rows = [values]
|
1016 |
|
|
for i, row in enumerate(rows):
|
1017 |
|
|
rows[i] = map(remove_col_rename, map(as_Value, row))
|
1018 |
2737
|
aaronmk
|
|
1019 |
2739
|
aaronmk
|
self.rows = rows
|
1020 |
2737
|
aaronmk
|
|
1021 |
|
|
def to_str(self, db):
|
1022 |
3520
|
aaronmk
|
return 'VALUES '+(', '.join((Tuple(*r).to_str(db) for r in self.rows)))
|
1023 |
2737
|
aaronmk
|
|
1024 |
2740
|
aaronmk
|
def NamedValues(name, cols, values):
|
1025 |
2745
|
aaronmk
|
'''
|
1026 |
3048
|
aaronmk
|
@param cols None|[...]
|
1027 |
2745
|
aaronmk
|
@post `cols` will be changed to Col objects with the table set to `name`.
|
1028 |
|
|
'''
|
1029 |
2834
|
aaronmk
|
table = NamedTable(name, Values(values), cols)
|
1030 |
3048
|
aaronmk
|
if cols != None: set_cols_table(table, cols)
|
1031 |
2834
|
aaronmk
|
return table
|
1032 |
2740
|
aaronmk
|
|
1033 |
2674
|
aaronmk
|
##### Database structure
|
1034 |
|
|
|
1035 |
4406
|
aaronmk
|
def is_text_col(db, col): return db.col_info(col).type == 'text'
|
1036 |
|
|
|
1037 |
2840
|
aaronmk
|
ensure_not_null_excs = (NoUnderlyingTableException, KeyError)
|
1038 |
|
|
|
1039 |
2851
|
aaronmk
|
def ensure_not_null(db, col, type_=None):
|
1040 |
2840
|
aaronmk
|
'''
|
1041 |
2855
|
aaronmk
|
@param col If type_ is not set, must have an underlying column.
|
1042 |
2851
|
aaronmk
|
@param type_ If set, overrides the underlying column's type.
|
1043 |
2840
|
aaronmk
|
@return EnsureNotNull|Col
|
1044 |
|
|
@throws ensure_not_null_excs
|
1045 |
|
|
'''
|
1046 |
2855
|
aaronmk
|
nullable = True
|
1047 |
|
|
try: typed_col = db.col_info(underlying_col(col))
|
1048 |
|
|
except NoUnderlyingTableException:
|
1049 |
3355
|
aaronmk
|
col = remove_col_rename(col)
|
1050 |
3429
|
aaronmk
|
if is_literal(col) and not is_null(col): nullable = False
|
1051 |
3355
|
aaronmk
|
elif type_ == None: raise
|
1052 |
2855
|
aaronmk
|
else:
|
1053 |
|
|
if type_ == None: type_ = typed_col.type
|
1054 |
|
|
nullable = typed_col.nullable
|
1055 |
|
|
|
1056 |
2953
|
aaronmk
|
if nullable:
|
1057 |
|
|
try: col = EnsureNotNull(col, type_)
|
1058 |
|
|
except KeyError, e:
|
1059 |
|
|
# Warn of no null sentinel for type, even if caller catches error
|
1060 |
|
|
warnings.warn(UserWarning(exc.str_(e)))
|
1061 |
|
|
raise
|
1062 |
|
|
|
1063 |
2840
|
aaronmk
|
return col
|
1064 |
3536
|
aaronmk
|
|
1065 |
|
|
def try_mk_not_null(db, value):
|
1066 |
|
|
'''
|
1067 |
|
|
Warning: This function does not guarantee that its result is NOT NULL.
|
1068 |
|
|
'''
|
1069 |
|
|
try: return ensure_not_null(db, value)
|
1070 |
|
|
except ensure_not_null_excs: return value
|