Project

General

Profile

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