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