Project

General

Profile

1 2211 aaronmk
# SQL code generation
2
3 2748 aaronmk
import copy
4 2276 aaronmk
import operator
5 2568 aaronmk
import re
6 2653 aaronmk
import UserDict
7 2953 aaronmk
import warnings
8 2276 aaronmk
9 2667 aaronmk
import dicts
10 2953 aaronmk
import exc
11 2701 aaronmk
import iters
12
import lists
13 2360 aaronmk
import objects
14 2222 aaronmk
import strings
15 2227 aaronmk
import util
16 2211 aaronmk
17 2587 aaronmk
##### Names
18 2499 aaronmk
19 2608 aaronmk
identifier_max_len = 63 # works for both PostgreSQL and MySQL
20 2587 aaronmk
21 2932 aaronmk
def concat(str_, suffix):
22 2609 aaronmk
    '''Preserves version so that it won't be truncated off the string, leading
23
    to collisions.'''
24 2613 aaronmk
    # Preserve version
25 2995 aaronmk
    match = re.match(r'^(.*?)((?:(?:#\d+)?\)?)*(?:\.\w+)?(?:::[\w ]+)*)$', str_)
26 2985 aaronmk
    if match:
27
        str_, old_suffix = match.groups()
28
        suffix = old_suffix+suffix
29 2613 aaronmk
30 2932 aaronmk
    return strings.concat(str_, suffix, identifier_max_len)
31 2587 aaronmk
32 2932 aaronmk
def truncate(str_): return concat(str_, '')
33 2842 aaronmk
34 2575 aaronmk
def is_safe_name(name):
35 2583 aaronmk
    '''A name is safe *and unambiguous* if it:
36
    * contains only *lowercase* word (\w) characters
37
    * doesn't start with a digit
38
    * contains "_", so that it's not a keyword
39 2984 aaronmk
    '''
40
    return re.match(r'^(?=.*_)(?!\d)[^\WA-Z]+$', name)
41 2568 aaronmk
42 2499 aaronmk
def esc_name(name, quote='"'):
43
    return quote + name.replace(quote, quote+quote) + quote
44
        # doubling an embedded quote escapes it in both PostgreSQL and MySQL
45
46 2513 aaronmk
def clean_name(name): return name.replace('"', '').replace('`', '')
47
48 3041 aaronmk
def esc_comment(comment): return '/*'+comment.replace('*/', '* /')+'*/'
49
50 2659 aaronmk
##### General SQL code objects
51 2219 aaronmk
52 2349 aaronmk
class MockDb:
53 2503 aaronmk
    def esc_value(self, value): return strings.repr_no_u(value)
54 2349 aaronmk
55 2499 aaronmk
    def esc_name(self, name): return esc_name(name)
56 2859 aaronmk
57
    def col_info(self, col):
58
        return TypedCol(col.name, '<type>', CustomCode('<default>'), True)
59
60 2349 aaronmk
mockDb = MockDb()
61
62 2514 aaronmk
class BasicObject(objects.BasicObject):
63
    def __init__(self, value): self.value = value
64
65
    def __str__(self): return clean_name(strings.repr_no_u(self))
66
67 2659 aaronmk
##### Unparameterized code objects
68
69 2514 aaronmk
class Code(BasicObject):
70 2658 aaronmk
    def to_str(self, db): raise NotImplementedError()
71 2349 aaronmk
72 2514 aaronmk
    def __repr__(self): return self.to_str(mockDb)
73 2211 aaronmk
74 2269 aaronmk
class CustomCode(Code):
75 2256 aaronmk
    def __init__(self, str_): self.str_ = str_
76
77
    def to_str(self, db): return self.str_
78
79 2815 aaronmk
def as_Code(value, db=None):
80
    '''
81
    @param db If set, runs db.std_code() on the value.
82
    '''
83
    if util.is_str(value):
84
        if db != None: value = db.std_code(value)
85
        return CustomCode(value)
86 2659 aaronmk
    else: return Literal(value)
87
88 2540 aaronmk
class Expr(Code):
89
    def __init__(self, expr): self.expr = expr
90
91
    def to_str(self, db): return '('+self.expr.to_str(db)+')'
92
93 3086 aaronmk
##### Names
94
95
class Name(Code):
96 3087 aaronmk
    def __init__(self, name):
97
        name = truncate(name)
98
99
        self.name = name
100 3086 aaronmk
101
    def to_str(self, db): return db.esc_name(self.name)
102
103
def as_Name(value):
104
    if isinstance(value, Code): return value
105
    else: return Name(value)
106
107 2335 aaronmk
##### Literal values
108
109 2216 aaronmk
class Literal(Code):
110 2211 aaronmk
    def __init__(self, value): self.value = value
111 2213 aaronmk
112
    def to_str(self, db): return db.esc_value(self.value)
113 2211 aaronmk
114 2400 aaronmk
def as_Value(value):
115
    if isinstance(value, Code): return value
116
    else: return Literal(value)
117
118 2216 aaronmk
def is_null(value): return isinstance(value, Literal) and value.value == None
119
120 2711 aaronmk
##### Derived elements
121
122
src_self = object() # tells Col that it is its own source column
123
124
class Derived(Code):
125
    def __init__(self, srcs):
126 2712 aaronmk
        '''An element which was derived from some other element(s).
127 2711 aaronmk
        @param srcs See self.set_srcs()
128
        '''
129
        self.set_srcs(srcs)
130
131 2713 aaronmk
    def set_srcs(self, srcs, overwrite=True):
132 2711 aaronmk
        '''
133
        @param srcs (self_type...)|src_self The element(s) this is derived from
134
        '''
135 2713 aaronmk
        if not overwrite and self.srcs != (): return # already set
136
137 2711 aaronmk
        if srcs == src_self: srcs = (self,)
138
        srcs = tuple(srcs) # make Col hashable
139
        self.srcs = srcs
140
141
    def _compare_on(self):
142
        compare_on = self.__dict__.copy()
143
        del compare_on['srcs'] # ignore
144
        return compare_on
145
146
def cols_srcs(cols): return lists.uniqify(iters.flatten((v.srcs for v in cols)))
147
148 2335 aaronmk
##### Tables
149
150 2712 aaronmk
class Table(Derived):
151 2991 aaronmk
    def __init__(self, name, schema=None, srcs=(), is_temp=False):
152 2211 aaronmk
        '''
153
        @param schema str|None (for no schema)
154 2712 aaronmk
        @param srcs (Table...)|src_self See Derived.set_srcs()
155 2211 aaronmk
        '''
156 2712 aaronmk
        Derived.__init__(self, srcs)
157
158 3091 aaronmk
        if util.is_str(name): name = truncate(name)
159 2843 aaronmk
160 2211 aaronmk
        self.name = name
161
        self.schema = schema
162 2991 aaronmk
        self.is_temp = is_temp
163 3000 aaronmk
        self.index_cols = {}
164 2211 aaronmk
165 2348 aaronmk
    def to_str(self, db):
166
        str_ = ''
167 3088 aaronmk
        if self.schema != None: str_ += as_Name(self.schema).to_str(db)+'.'
168
        str_ += as_Name(self.name).to_str(db)
169 2348 aaronmk
        return str_
170 2336 aaronmk
171
    def to_Table(self): return self
172 3000 aaronmk
173
    def _compare_on(self):
174
        compare_on = Derived._compare_on(self)
175
        del compare_on['index_cols'] # ignore
176
        return compare_on
177 2211 aaronmk
178 2835 aaronmk
def is_underlying_table(table):
179
    return isinstance(table, Table) and table.to_Table() is table
180 2832 aaronmk
181 2902 aaronmk
class NoUnderlyingTableException(Exception): pass
182
183
def underlying_table(table):
184
    table = remove_table_rename(table)
185
    if not is_underlying_table(table): raise NoUnderlyingTableException
186
    return table
187
188 2776 aaronmk
def as_Table(table, schema=None):
189 2270 aaronmk
    if table == None or isinstance(table, Code): return table
190 2776 aaronmk
    else: return Table(table, schema)
191 2219 aaronmk
192 2707 aaronmk
def suffixed_table(table, suffix): return Table(table.name+suffix, table.schema)
193
194 2336 aaronmk
class NamedTable(Table):
195
    def __init__(self, name, code, cols=None):
196
        Table.__init__(self, name)
197
198 3016 aaronmk
        code = as_Table(code)
199 2741 aaronmk
        if not isinstance(code, (Table, FunctionCall, Expr)): code = Expr(code)
200 3020 aaronmk
        if cols != None: cols = [to_name_only_col(c).to_Col() for c in cols]
201 2336 aaronmk
202
        self.code = code
203
        self.cols = cols
204
205
    def to_str(self, db):
206 3026 aaronmk
        str_ = self.code.to_str(db)
207
        if str_.find('\n') >= 0: whitespace = '\n'
208
        else: whitespace = ' '
209
        str_ += whitespace+'AS '+Table.to_str(self, db)
210 2742 aaronmk
        if self.cols != None:
211
            str_ += ' ('+(', '.join((c.to_str(db) for c in self.cols)))+')'
212 2336 aaronmk
        return str_
213
214
    def to_Table(self): return Table(self.name)
215
216 2753 aaronmk
def remove_table_rename(table):
217
    if isinstance(table, NamedTable): table = table.code
218
    return table
219
220 2335 aaronmk
##### Columns
221
222 2711 aaronmk
class Col(Derived):
223 2701 aaronmk
    def __init__(self, name, table=None, srcs=()):
224 2211 aaronmk
        '''
225
        @param table Table|None (for no table)
226 2711 aaronmk
        @param srcs (Col...)|src_self See Derived.set_srcs()
227 2211 aaronmk
        '''
228 2711 aaronmk
        Derived.__init__(self, srcs)
229
230 3091 aaronmk
        if util.is_str(name): name = truncate(name)
231 2241 aaronmk
        if util.is_str(table): table = Table(table)
232 2211 aaronmk
        assert table == None or isinstance(table, Table)
233
234
        self.name = name
235
        self.table = table
236
237 2989 aaronmk
    def to_str(self, db, for_str=False):
238 3088 aaronmk
        str_ = as_Name(self.name).to_str(db)
239 2989 aaronmk
        if for_str: str_ = clean_name(str_)
240 2933 aaronmk
        if self.table != None:
241 2989 aaronmk
            table = self.table.to_Table()
242
            if for_str: str_ = concat(str(table), '.'+str_)
243
            else: str_ = table.to_str(db)+'.'+str_
244 2211 aaronmk
        return str_
245 2314 aaronmk
246 2989 aaronmk
    def __str__(self): return self.to_str(mockDb, for_str=True)
247 2933 aaronmk
248 2314 aaronmk
    def to_Col(self): return self
249 2211 aaronmk
250 2767 aaronmk
def is_table_col(col): return isinstance(col, Col) and col.table != None
251 2393 aaronmk
252 3000 aaronmk
def index_col(col):
253
    if not is_table_col(col): return None
254
    return col.table.index_cols.get(col.name, None)
255 2999 aaronmk
256 3024 aaronmk
def is_temp_col(col): return is_table_col(col) and col.table.is_temp
257 2996 aaronmk
258 2563 aaronmk
def as_Col(col, table=None, name=None):
259
    '''
260
    @param name If not None, any non-Col input will be renamed using NamedCol.
261
    '''
262
    if name != None:
263
        col = as_Value(col)
264
        if not isinstance(col, Col): col = NamedCol(name, col)
265 2333 aaronmk
266
    if isinstance(col, Code): return col
267 2260 aaronmk
    else: return Col(col, table)
268
269 3093 aaronmk
def with_table(col, table):
270
    col = copy.copy(col) # don't modify input!
271
    col.table = table
272
    return col
273
274 2750 aaronmk
def with_default_table(col, table, overwrite=False):
275 2747 aaronmk
    col = as_Col(col)
276 2750 aaronmk
    if not isinstance(col, NamedCol) and (overwrite or col.table == None):
277 3093 aaronmk
        col = with_table(col, table)
278 2747 aaronmk
    return col
279
280 2744 aaronmk
def set_cols_table(table, cols):
281
    table = as_Table(table)
282
283
    for i, col in enumerate(cols):
284
        col = cols[i] = as_Col(col)
285
        col.table = table
286
287 2401 aaronmk
def to_name_only_col(col, check_table=None):
288
    col = as_Col(col)
289 3020 aaronmk
    if not is_table_col(col): return col
290 2401 aaronmk
291
    if check_table != None:
292
        table = col.table
293
        assert table == None or table == check_table
294
    return Col(col.name)
295
296 2993 aaronmk
def suffixed_col(col, suffix):
297
    return Col(concat(col.name, suffix), col.table, col.srcs)
298
299 2323 aaronmk
class NamedCol(Col):
300 2229 aaronmk
    def __init__(self, name, code):
301 2310 aaronmk
        Col.__init__(self, name)
302
303 3016 aaronmk
        code = as_Value(code)
304 2229 aaronmk
305
        self.code = code
306
307
    def to_str(self, db):
308 2310 aaronmk
        return self.code.to_str(db)+' AS '+Col.to_str(self, db)
309 2314 aaronmk
310
    def to_Col(self): return Col(self.name)
311 2229 aaronmk
312 2462 aaronmk
def remove_col_rename(col):
313
    if isinstance(col, NamedCol): col = col.code
314
    return col
315
316 2830 aaronmk
def underlying_col(col):
317
    col = remove_col_rename(col)
318 2849 aaronmk
    if not isinstance(col, Col): raise NoUnderlyingTableException
319
320 2902 aaronmk
    return Col(col.name, underlying_table(col.table), col.srcs)
321 2830 aaronmk
322 2703 aaronmk
def wrap(wrap_func, value):
323
    '''Wraps a value, propagating any column renaming to the returned value.'''
324
    if isinstance(value, NamedCol):
325
        return NamedCol(value.name, wrap_func(value.code))
326
    else: return wrap_func(value)
327
328 2667 aaronmk
class ColDict(dicts.DictProxy):
329 2564 aaronmk
    '''A dict that automatically makes inserted entries Col objects'''
330
331 2645 aaronmk
    def __init__(self, db, keys_table, dict_={}):
332 2667 aaronmk
        dicts.DictProxy.__init__(self, {})
333
334 2645 aaronmk
        keys_table = as_Table(keys_table)
335
336 2642 aaronmk
        self.db = db
337 2641 aaronmk
        self.table = keys_table
338 2653 aaronmk
        self.update(dict_) # after setting vars because __setitem__() needs them
339 2641 aaronmk
340 2667 aaronmk
    def copy(self): return ColDict(self.db, self.table, self.inner.copy())
341 2655 aaronmk
342 2667 aaronmk
    def __getitem__(self, key):
343
        return dicts.DictProxy.__getitem__(self, self._key(key))
344 2653 aaronmk
345 2564 aaronmk
    def __setitem__(self, key, value):
346 2642 aaronmk
        key = self._key(key)
347 2819 aaronmk
        if value == None: value = self.db.col_info(key).default
348 2667 aaronmk
        dicts.DictProxy.__setitem__(self, key, as_Col(value, name=key.name))
349 2564 aaronmk
350 2641 aaronmk
    def _key(self, key): return as_Col(key, self.table)
351 2564 aaronmk
352 2524 aaronmk
##### Functions
353
354 2912 aaronmk
Function = Table
355 2911 aaronmk
as_Function = as_Table
356
357 2691 aaronmk
class InternalFunction(CustomCode): pass
358
359 2941 aaronmk
class NamedArg(NamedCol):
360
    def __init__(self, name, value):
361
        NamedCol.__init__(self, name, value)
362
363
    def to_str(self, db):
364
        return Col.to_str(self, db)+' := '+self.code.to_str(db)
365
366 2524 aaronmk
class FunctionCall(Code):
367 2941 aaronmk
    def __init__(self, function, *args, **kw_args):
368 2524 aaronmk
        '''
369 2690 aaronmk
        @param args [Code|literal-value...] The function's arguments
370 2524 aaronmk
        '''
371 3016 aaronmk
        function = as_Function(function)
372 2941 aaronmk
        def filter_(arg): return remove_col_rename(as_Value(arg))
373
        args = map(filter_, args)
374
        args += [NamedArg(k, filter_(v)) for k, v in kw_args.iteritems()]
375 2524 aaronmk
376
        self.function = function
377
        self.args = args
378
379
    def to_str(self, db):
380
        args_str = ', '.join((v.to_str(db) for v in self.args))
381
        return self.function.to_str(db)+'('+args_str+')'
382
383 2533 aaronmk
def wrap_in_func(function, value):
384
    '''Wraps a value inside a function call.
385
    Propagates any column renaming to the returned value.
386
    '''
387 2703 aaronmk
    return wrap(lambda v: FunctionCall(function, v), value)
388 2533 aaronmk
389 2561 aaronmk
def unwrap_func_call(func_call, check_name=None):
390
    '''Unwraps any function call to its first argument.
391
    Also removes any column renaming.
392
    '''
393
    func_call = remove_col_rename(func_call)
394
    if not isinstance(func_call, FunctionCall): return func_call
395
396
    if check_name != None:
397
        name = func_call.function.name
398
        assert name == None or name == check_name
399
    return func_call.args[0]
400
401 2986 aaronmk
##### Casts
402
403
class Cast(FunctionCall):
404
    def __init__(self, type_, value):
405
        value = as_Value(value)
406
407
        self.type_ = type_
408
        self.value = value
409
410
    def to_str(self, db):
411
        return 'CAST('+self.value.to_str(db)+' AS '+self.type_+')'
412
413 2335 aaronmk
##### Conditions
414 2259 aaronmk
415 2398 aaronmk
class ColValueCond(Code):
416
    def __init__(self, col, value):
417
        value = as_ValueCond(value)
418
419
        self.col = col
420
        self.value = value
421
422
    def to_str(self, db): return self.value.to_str(db, self.col)
423
424 2577 aaronmk
def combine_conds(conds, keyword=None):
425
    '''
426
    @param keyword The keyword to add before the conditions, if any
427
    '''
428
    str_ = ''
429
    if keyword != None:
430
        if conds == []: whitespace = ''
431
        elif len(conds) == 1: whitespace = ' '
432
        else: whitespace = '\n'
433
        str_ += keyword+whitespace
434
435
    str_ += '\nAND '.join(conds)
436
    return str_
437
438 2398 aaronmk
##### Condition column comparisons
439
440 2514 aaronmk
class ValueCond(BasicObject):
441 2213 aaronmk
    def __init__(self, value):
442 2858 aaronmk
        value = remove_col_rename(as_Value(value))
443 2213 aaronmk
444
        self.value = value
445 2214 aaronmk
446 2216 aaronmk
    def to_str(self, db, left_value):
447 2214 aaronmk
        '''
448 2216 aaronmk
        @param left_value The Code object that the condition is being applied on
449 2214 aaronmk
        '''
450
        raise NotImplemented()
451 2228 aaronmk
452 2514 aaronmk
    def __repr__(self): return self.to_str(mockDb, '<left_value>')
453 2211 aaronmk
454
class CompareCond(ValueCond):
455
    def __init__(self, value, operator='='):
456 2222 aaronmk
        '''
457
        @param operator By default, compares NULL values literally. Use '~=' or
458
            '~!=' to pass NULLs through.
459
        '''
460 2211 aaronmk
        ValueCond.__init__(self, value)
461
        self.operator = operator
462
463 2216 aaronmk
    def to_str(self, db, left_value):
464 2858 aaronmk
        left_value = remove_col_rename(as_Col(left_value))
465 2216 aaronmk
466 2222 aaronmk
        right_value = self.value
467
468
        # Parse operator
469 2216 aaronmk
        operator = self.operator
470 2222 aaronmk
        passthru_null_ref = [False]
471
        operator = strings.remove_prefix('~', operator, passthru_null_ref)
472
        neg_ref = [False]
473
        operator = strings.remove_prefix('!', operator, neg_ref)
474 2844 aaronmk
        equals = operator.endswith('=') # also includes <=, >=
475 2222 aaronmk
476 2825 aaronmk
        # Handle nullable columns
477
        check_null = False
478 2844 aaronmk
        if not passthru_null_ref[0]: # NULLs compare equal
479 2857 aaronmk
            try: left_value = ensure_not_null(db, left_value)
480 2844 aaronmk
            except ensure_not_null_excs: # fall back to alternate method
481
                check_null = equals and isinstance(right_value, Col)
482 2837 aaronmk
            else:
483 2857 aaronmk
                if isinstance(left_value, EnsureNotNull):
484
                    right_value = ensure_not_null(db, right_value,
485
                        left_value.type) # apply same function to both sides
486 2825 aaronmk
487 2844 aaronmk
        if equals and is_null(right_value): operator = 'IS'
488
489 2825 aaronmk
        left = left_value.to_str(db)
490
        right = right_value.to_str(db)
491
492 2222 aaronmk
        # Create str
493
        str_ = left+' '+operator+' '+right
494 2825 aaronmk
        if check_null:
495 2578 aaronmk
            str_ = '('+str_+' OR ('+left+' IS NULL AND '+right+' IS NULL))'
496
        if neg_ref[0]: str_ = 'NOT '+str_
497 2222 aaronmk
        return str_
498 2216 aaronmk
499 2260 aaronmk
# Tells as_ValueCond() to assume a non-ValueCond is a literal value
500
assume_literal = object()
501
502
def as_ValueCond(value, default_table=assume_literal):
503
    if not isinstance(value, ValueCond):
504
        if default_table is not assume_literal:
505 2748 aaronmk
            value = with_default_table(value, default_table)
506 2260 aaronmk
        return CompareCond(value)
507 2216 aaronmk
    else: return value
508 2219 aaronmk
509 2335 aaronmk
##### Joins
510
511 2352 aaronmk
join_same = object() # tells Join the left and right columns have the same name
512 2260 aaronmk
513 2353 aaronmk
# Tells Join the left and right columns have the same name and are never NULL
514
join_same_not_null = object()
515
516 2260 aaronmk
filter_out = object() # tells Join to filter out rows that match the join
517
518 2514 aaronmk
class Join(BasicObject):
519 2746 aaronmk
    def __init__(self, table, mapping={}, type_=None):
520 2260 aaronmk
        '''
521
        @param mapping dict(right_table_col=left_table_col, ...)
522 2352 aaronmk
            * if left_table_col is join_same: left_table_col = right_table_col
523 2353 aaronmk
              * Note that right_table_col must be a string
524
            * if left_table_col is join_same_not_null:
525
              left_table_col = right_table_col and both have NOT NULL constraint
526
              * Note that right_table_col must be a string
527 2260 aaronmk
        @param type_ None (for plain join)|str (e.g. 'LEFT')|filter_out
528
            * filter_out: equivalent to 'LEFT' with the query filtered by
529
              `table_pkey IS NULL` (indicating no match)
530
        '''
531
        if util.is_str(table): table = Table(table)
532
        assert type_ == None or util.is_str(type_) or type_ is filter_out
533
534
        self.table = table
535
        self.mapping = mapping
536
        self.type_ = type_
537
538 2749 aaronmk
    def to_str(self, db, left_table_):
539 2260 aaronmk
        def join(entry):
540
            '''Parses non-USING joins'''
541
            right_table_col, left_table_col = entry
542
543 2353 aaronmk
            # Switch order (right_table_col is on the left in the comparison)
544
            left = right_table_col
545
            right = left_table_col
546 2749 aaronmk
            left_table = self.table
547
            right_table = left_table_
548 2353 aaronmk
549 2747 aaronmk
            # Parse left side
550 2748 aaronmk
            left = with_default_table(left, left_table)
551 2747 aaronmk
552 2260 aaronmk
            # Parse special values
553 2747 aaronmk
            left_on_right = Col(left.name, right_table)
554
            if right is join_same: right = left_on_right
555 2353 aaronmk
            elif right is join_same_not_null:
556 2747 aaronmk
                right = CompareCond(left_on_right, '~=')
557 2260 aaronmk
558 2747 aaronmk
            # Parse right side
559 2353 aaronmk
            right = as_ValueCond(right, right_table)
560 2747 aaronmk
561
            return right.to_str(db, left)
562 2260 aaronmk
563 2265 aaronmk
        # Create join condition
564
        type_ = self.type_
565 2276 aaronmk
        joins = self.mapping
566 2746 aaronmk
        if joins == {}: join_cond = None
567
        elif type_ is not filter_out and reduce(operator.and_,
568 2460 aaronmk
            (v is join_same_not_null for v in joins.itervalues())):
569 2260 aaronmk
            # all cols w/ USING, so can use simpler USING syntax
570 2747 aaronmk
            cols = map(to_name_only_col, joins.iterkeys())
571
            join_cond = 'USING ('+(', '.join((c.to_str(db) for c in cols)))+')'
572 2757 aaronmk
        else: join_cond = combine_conds(map(join, joins.iteritems()), 'ON')
573 2260 aaronmk
574 2757 aaronmk
        if isinstance(self.table, NamedTable): whitespace = '\n'
575
        else: whitespace = ' '
576
577 2260 aaronmk
        # Create join
578
        if type_ is filter_out: type_ = 'LEFT'
579 2266 aaronmk
        str_ = ''
580
        if type_ != None: str_ += type_+' '
581 2757 aaronmk
        str_ += 'JOIN'+whitespace+self.table.to_str(db)
582
        if join_cond != None: str_ += whitespace+join_cond
583 2266 aaronmk
        return str_
584 2349 aaronmk
585 2514 aaronmk
    def __repr__(self): return self.to_str(mockDb, '<left_table>')
586 2424 aaronmk
587
##### Value exprs
588
589 3089 aaronmk
all_cols = CustomCode('*')
590
591 2737 aaronmk
default = CustomCode('DEFAULT')
592
593 3090 aaronmk
row_count = FunctionCall(InternalFunction('COUNT'), all_cols)
594 2674 aaronmk
595 3061 aaronmk
class Coalesce(FunctionCall):
596
    def __init__(self, *args):
597
        FunctionCall.__init__(self, InternalFunction('COALESCE'), *args)
598 3060 aaronmk
599 3062 aaronmk
class Nullif(FunctionCall):
600
    def __init__(self, *args):
601
        FunctionCall.__init__(self, InternalFunction('NULLIF'), *args)
602
603 2850 aaronmk
# See <http://www.postgresql.org/docs/8.3/static/datatype-numeric.html>
604 2958 aaronmk
null_sentinels = {
605
    'character varying': r'\N',
606
    'double precision': 'NaN',
607
    'integer': 2147483647,
608
    'text': r'\N',
609
    'timestamp with time zone': 'infinity'
610
}
611 2692 aaronmk
612 3061 aaronmk
class EnsureNotNull(Coalesce):
613 2850 aaronmk
    def __init__(self, value, type_):
614 3061 aaronmk
        Coalesce.__init__(self, as_Col(value),
615 2988 aaronmk
            Cast(type_, null_sentinels[type_]))
616 2850 aaronmk
617
        self.type = type_
618 3001 aaronmk
619
    def to_str(self, db):
620
        col = self.args[0]
621
        index_col_ = index_col(col)
622
        if index_col_ != None: return index_col_.to_str(db)
623 3061 aaronmk
        return Coalesce.to_str(self, db)
624 2850 aaronmk
625 2737 aaronmk
##### Table exprs
626
627
class Values(Code):
628
    def __init__(self, values):
629 2739 aaronmk
        '''
630
        @param values [...]|[[...], ...] Can be one or multiple rows.
631
        '''
632
        rows = values
633
        if len(values) >= 1 and not lists.is_seq(values[0]): # only one row
634
            rows = [values]
635
        for i, row in enumerate(rows):
636
            rows[i] = map(remove_col_rename, map(as_Value, row))
637 2737 aaronmk
638 2739 aaronmk
        self.rows = rows
639 2737 aaronmk
640
    def to_str(self, db):
641 2739 aaronmk
        def row_str(row):
642
            return '('+(', '.join((v.to_str(db) for v in row)))+')'
643
        return 'VALUES '+(', '.join(map(row_str, self.rows)))
644 2737 aaronmk
645 2740 aaronmk
def NamedValues(name, cols, values):
646 2745 aaronmk
    '''
647 3048 aaronmk
    @param cols None|[...]
648 2745 aaronmk
    @post `cols` will be changed to Col objects with the table set to `name`.
649
    '''
650 2834 aaronmk
    table = NamedTable(name, Values(values), cols)
651 3048 aaronmk
    if cols != None: set_cols_table(table, cols)
652 2834 aaronmk
    return table
653 2740 aaronmk
654 2674 aaronmk
##### Database structure
655
656
class TypedCol(Col):
657 2871 aaronmk
    def __init__(self, name, type_, default=None, nullable=True,
658
        constraints=None):
659 2818 aaronmk
        assert default == None or isinstance(default, Code)
660
661 2674 aaronmk
        Col.__init__(self, name)
662
663
        self.type = type_
664 2818 aaronmk
        self.default = default
665
        self.nullable = nullable
666 2871 aaronmk
        self.constraints = constraints
667 2674 aaronmk
668 2818 aaronmk
    def to_str(self, db):
669
        str_ = Col.to_str(self, db)+' '+self.type
670
        if not self.nullable: str_ += ' NOT NULL'
671
        if self.default != None: str_ += ' DEFAULT '+self.default.to_str(db)
672 2871 aaronmk
        if self.constraints != None: str_ += ' '+self.constraints
673 2818 aaronmk
        return str_
674 2674 aaronmk
675
    def to_Col(self): return Col(self.name)
676 2822 aaronmk
677 2840 aaronmk
ensure_not_null_excs = (NoUnderlyingTableException, KeyError)
678
679 2851 aaronmk
def ensure_not_null(db, col, type_=None):
680 2840 aaronmk
    '''
681 2855 aaronmk
    @param col If type_ is not set, must have an underlying column.
682 2851 aaronmk
    @param type_ If set, overrides the underlying column's type.
683 2840 aaronmk
    @return EnsureNotNull|Col
684
    @throws ensure_not_null_excs
685
    '''
686 2855 aaronmk
    nullable = True
687
    try: typed_col = db.col_info(underlying_col(col))
688
    except NoUnderlyingTableException:
689
        if type_ == None: raise
690
    else:
691
        if type_ == None: type_ = typed_col.type
692
        nullable = typed_col.nullable
693
694 2953 aaronmk
    if nullable:
695
        try: col = EnsureNotNull(col, type_)
696
        except KeyError, e:
697
            # Warn of no null sentinel for type, even if caller catches error
698
            warnings.warn(UserWarning(exc.str_(e)))
699
            raise
700
701 2840 aaronmk
    return col