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 3101 aaronmk
def suffixed_table(table, suffix):
193
    return Table(concat(table.name, suffix), table.schema)
194 2707 aaronmk
195 2336 aaronmk
class NamedTable(Table):
196
    def __init__(self, name, code, cols=None):
197
        Table.__init__(self, name)
198
199 3016 aaronmk
        code = as_Table(code)
200 2741 aaronmk
        if not isinstance(code, (Table, FunctionCall, Expr)): code = Expr(code)
201 3020 aaronmk
        if cols != None: cols = [to_name_only_col(c).to_Col() for c in cols]
202 2336 aaronmk
203
        self.code = code
204
        self.cols = cols
205
206
    def to_str(self, db):
207 3026 aaronmk
        str_ = self.code.to_str(db)
208
        if str_.find('\n') >= 0: whitespace = '\n'
209
        else: whitespace = ' '
210
        str_ += whitespace+'AS '+Table.to_str(self, db)
211 2742 aaronmk
        if self.cols != None:
212
            str_ += ' ('+(', '.join((c.to_str(db) for c in self.cols)))+')'
213 2336 aaronmk
        return str_
214
215
    def to_Table(self): return Table(self.name)
216
217 2753 aaronmk
def remove_table_rename(table):
218
    if isinstance(table, NamedTable): table = table.code
219
    return table
220
221 2335 aaronmk
##### Columns
222
223 2711 aaronmk
class Col(Derived):
224 2701 aaronmk
    def __init__(self, name, table=None, srcs=()):
225 2211 aaronmk
        '''
226
        @param table Table|None (for no table)
227 2711 aaronmk
        @param srcs (Col...)|src_self See Derived.set_srcs()
228 2211 aaronmk
        '''
229 2711 aaronmk
        Derived.__init__(self, srcs)
230
231 3091 aaronmk
        if util.is_str(name): name = truncate(name)
232 2241 aaronmk
        if util.is_str(table): table = Table(table)
233 2211 aaronmk
        assert table == None or isinstance(table, Table)
234
235
        self.name = name
236
        self.table = table
237
238 2989 aaronmk
    def to_str(self, db, for_str=False):
239 3088 aaronmk
        str_ = as_Name(self.name).to_str(db)
240 2989 aaronmk
        if for_str: str_ = clean_name(str_)
241 2933 aaronmk
        if self.table != None:
242 2989 aaronmk
            table = self.table.to_Table()
243
            if for_str: str_ = concat(str(table), '.'+str_)
244
            else: str_ = table.to_str(db)+'.'+str_
245 2211 aaronmk
        return str_
246 2314 aaronmk
247 2989 aaronmk
    def __str__(self): return self.to_str(mockDb, for_str=True)
248 2933 aaronmk
249 2314 aaronmk
    def to_Col(self): return self
250 2211 aaronmk
251 2767 aaronmk
def is_table_col(col): return isinstance(col, Col) and col.table != None
252 2393 aaronmk
253 3000 aaronmk
def index_col(col):
254
    if not is_table_col(col): return None
255 3104 aaronmk
256
    table = col.table
257
    try: name = table.index_cols[col.name]
258
    except KeyError: return None
259
    else: return Col(name, table, col.srcs)
260 2999 aaronmk
261 3024 aaronmk
def is_temp_col(col): return is_table_col(col) and col.table.is_temp
262 2996 aaronmk
263 2563 aaronmk
def as_Col(col, table=None, name=None):
264
    '''
265
    @param name If not None, any non-Col input will be renamed using NamedCol.
266
    '''
267
    if name != None:
268
        col = as_Value(col)
269
        if not isinstance(col, Col): col = NamedCol(name, col)
270 2333 aaronmk
271
    if isinstance(col, Code): return col
272 2260 aaronmk
    else: return Col(col, table)
273
274 3093 aaronmk
def with_table(col, table):
275 3098 aaronmk
    if not isinstance(col, NamedCol):
276
        col = copy.copy(col) # don't modify input!
277
        col.table = table
278 3093 aaronmk
    return col
279
280 3100 aaronmk
def with_default_table(col, table):
281 2747 aaronmk
    col = as_Col(col)
282 3100 aaronmk
    if col.table == None: col = with_table(col, table)
283 2747 aaronmk
    return col
284
285 2744 aaronmk
def set_cols_table(table, cols):
286
    table = as_Table(table)
287
288
    for i, col in enumerate(cols):
289
        col = cols[i] = as_Col(col)
290
        col.table = table
291
292 2401 aaronmk
def to_name_only_col(col, check_table=None):
293
    col = as_Col(col)
294 3020 aaronmk
    if not is_table_col(col): return col
295 2401 aaronmk
296
    if check_table != None:
297
        table = col.table
298
        assert table == None or table == check_table
299
    return Col(col.name)
300
301 2993 aaronmk
def suffixed_col(col, suffix):
302
    return Col(concat(col.name, suffix), col.table, col.srcs)
303
304 2323 aaronmk
class NamedCol(Col):
305 2229 aaronmk
    def __init__(self, name, code):
306 2310 aaronmk
        Col.__init__(self, name)
307
308 3016 aaronmk
        code = as_Value(code)
309 2229 aaronmk
310
        self.code = code
311
312
    def to_str(self, db):
313 2310 aaronmk
        return self.code.to_str(db)+' AS '+Col.to_str(self, db)
314 2314 aaronmk
315
    def to_Col(self): return Col(self.name)
316 2229 aaronmk
317 2462 aaronmk
def remove_col_rename(col):
318
    if isinstance(col, NamedCol): col = col.code
319
    return col
320
321 2830 aaronmk
def underlying_col(col):
322
    col = remove_col_rename(col)
323 2849 aaronmk
    if not isinstance(col, Col): raise NoUnderlyingTableException
324
325 2902 aaronmk
    return Col(col.name, underlying_table(col.table), col.srcs)
326 2830 aaronmk
327 2703 aaronmk
def wrap(wrap_func, value):
328
    '''Wraps a value, propagating any column renaming to the returned value.'''
329
    if isinstance(value, NamedCol):
330
        return NamedCol(value.name, wrap_func(value.code))
331
    else: return wrap_func(value)
332
333 2667 aaronmk
class ColDict(dicts.DictProxy):
334 2564 aaronmk
    '''A dict that automatically makes inserted entries Col objects'''
335
336 2645 aaronmk
    def __init__(self, db, keys_table, dict_={}):
337 2667 aaronmk
        dicts.DictProxy.__init__(self, {})
338
339 2645 aaronmk
        keys_table = as_Table(keys_table)
340
341 2642 aaronmk
        self.db = db
342 2641 aaronmk
        self.table = keys_table
343 2653 aaronmk
        self.update(dict_) # after setting vars because __setitem__() needs them
344 2641 aaronmk
345 2667 aaronmk
    def copy(self): return ColDict(self.db, self.table, self.inner.copy())
346 2655 aaronmk
347 2667 aaronmk
    def __getitem__(self, key):
348
        return dicts.DictProxy.__getitem__(self, self._key(key))
349 2653 aaronmk
350 2564 aaronmk
    def __setitem__(self, key, value):
351 2642 aaronmk
        key = self._key(key)
352 2819 aaronmk
        if value == None: value = self.db.col_info(key).default
353 2667 aaronmk
        dicts.DictProxy.__setitem__(self, key, as_Col(value, name=key.name))
354 2564 aaronmk
355 2641 aaronmk
    def _key(self, key): return as_Col(key, self.table)
356 2564 aaronmk
357 2524 aaronmk
##### Functions
358
359 2912 aaronmk
Function = Table
360 2911 aaronmk
as_Function = as_Table
361
362 2691 aaronmk
class InternalFunction(CustomCode): pass
363
364 2941 aaronmk
class NamedArg(NamedCol):
365
    def __init__(self, name, value):
366
        NamedCol.__init__(self, name, value)
367
368
    def to_str(self, db):
369
        return Col.to_str(self, db)+' := '+self.code.to_str(db)
370
371 2524 aaronmk
class FunctionCall(Code):
372 2941 aaronmk
    def __init__(self, function, *args, **kw_args):
373 2524 aaronmk
        '''
374 2690 aaronmk
        @param args [Code|literal-value...] The function's arguments
375 2524 aaronmk
        '''
376 3016 aaronmk
        function = as_Function(function)
377 2941 aaronmk
        def filter_(arg): return remove_col_rename(as_Value(arg))
378
        args = map(filter_, args)
379
        args += [NamedArg(k, filter_(v)) for k, v in kw_args.iteritems()]
380 2524 aaronmk
381
        self.function = function
382
        self.args = args
383
384
    def to_str(self, db):
385
        args_str = ', '.join((v.to_str(db) for v in self.args))
386
        return self.function.to_str(db)+'('+args_str+')'
387
388 2533 aaronmk
def wrap_in_func(function, value):
389
    '''Wraps a value inside a function call.
390
    Propagates any column renaming to the returned value.
391
    '''
392 2703 aaronmk
    return wrap(lambda v: FunctionCall(function, v), value)
393 2533 aaronmk
394 2561 aaronmk
def unwrap_func_call(func_call, check_name=None):
395
    '''Unwraps any function call to its first argument.
396
    Also removes any column renaming.
397
    '''
398
    func_call = remove_col_rename(func_call)
399
    if not isinstance(func_call, FunctionCall): return func_call
400
401
    if check_name != None:
402
        name = func_call.function.name
403
        assert name == None or name == check_name
404
    return func_call.args[0]
405
406 2986 aaronmk
##### Casts
407
408
class Cast(FunctionCall):
409
    def __init__(self, type_, value):
410
        value = as_Value(value)
411
412
        self.type_ = type_
413
        self.value = value
414
415
    def to_str(self, db):
416
        return 'CAST('+self.value.to_str(db)+' AS '+self.type_+')'
417
418 2335 aaronmk
##### Conditions
419 2259 aaronmk
420 2398 aaronmk
class ColValueCond(Code):
421
    def __init__(self, col, value):
422
        value = as_ValueCond(value)
423
424
        self.col = col
425
        self.value = value
426
427
    def to_str(self, db): return self.value.to_str(db, self.col)
428
429 2577 aaronmk
def combine_conds(conds, keyword=None):
430
    '''
431
    @param keyword The keyword to add before the conditions, if any
432
    '''
433
    str_ = ''
434
    if keyword != None:
435
        if conds == []: whitespace = ''
436
        elif len(conds) == 1: whitespace = ' '
437
        else: whitespace = '\n'
438
        str_ += keyword+whitespace
439
440
    str_ += '\nAND '.join(conds)
441
    return str_
442
443 2398 aaronmk
##### Condition column comparisons
444
445 2514 aaronmk
class ValueCond(BasicObject):
446 2213 aaronmk
    def __init__(self, value):
447 2858 aaronmk
        value = remove_col_rename(as_Value(value))
448 2213 aaronmk
449
        self.value = value
450 2214 aaronmk
451 2216 aaronmk
    def to_str(self, db, left_value):
452 2214 aaronmk
        '''
453 2216 aaronmk
        @param left_value The Code object that the condition is being applied on
454 2214 aaronmk
        '''
455
        raise NotImplemented()
456 2228 aaronmk
457 2514 aaronmk
    def __repr__(self): return self.to_str(mockDb, '<left_value>')
458 2211 aaronmk
459
class CompareCond(ValueCond):
460
    def __init__(self, value, operator='='):
461 2222 aaronmk
        '''
462
        @param operator By default, compares NULL values literally. Use '~=' or
463
            '~!=' to pass NULLs through.
464
        '''
465 2211 aaronmk
        ValueCond.__init__(self, value)
466
        self.operator = operator
467
468 2216 aaronmk
    def to_str(self, db, left_value):
469 2858 aaronmk
        left_value = remove_col_rename(as_Col(left_value))
470 2216 aaronmk
471 2222 aaronmk
        right_value = self.value
472
473
        # Parse operator
474 2216 aaronmk
        operator = self.operator
475 2222 aaronmk
        passthru_null_ref = [False]
476
        operator = strings.remove_prefix('~', operator, passthru_null_ref)
477
        neg_ref = [False]
478
        operator = strings.remove_prefix('!', operator, neg_ref)
479 2844 aaronmk
        equals = operator.endswith('=') # also includes <=, >=
480 2222 aaronmk
481 2825 aaronmk
        # Handle nullable columns
482
        check_null = False
483 2844 aaronmk
        if not passthru_null_ref[0]: # NULLs compare equal
484 2857 aaronmk
            try: left_value = ensure_not_null(db, left_value)
485 2844 aaronmk
            except ensure_not_null_excs: # fall back to alternate method
486
                check_null = equals and isinstance(right_value, Col)
487 2837 aaronmk
            else:
488 2857 aaronmk
                if isinstance(left_value, EnsureNotNull):
489
                    right_value = ensure_not_null(db, right_value,
490
                        left_value.type) # apply same function to both sides
491 2825 aaronmk
492 2844 aaronmk
        if equals and is_null(right_value): operator = 'IS'
493
494 2825 aaronmk
        left = left_value.to_str(db)
495
        right = right_value.to_str(db)
496
497 2222 aaronmk
        # Create str
498
        str_ = left+' '+operator+' '+right
499 2825 aaronmk
        if check_null:
500 2578 aaronmk
            str_ = '('+str_+' OR ('+left+' IS NULL AND '+right+' IS NULL))'
501
        if neg_ref[0]: str_ = 'NOT '+str_
502 2222 aaronmk
        return str_
503 2216 aaronmk
504 2260 aaronmk
# Tells as_ValueCond() to assume a non-ValueCond is a literal value
505
assume_literal = object()
506
507
def as_ValueCond(value, default_table=assume_literal):
508
    if not isinstance(value, ValueCond):
509
        if default_table is not assume_literal:
510 2748 aaronmk
            value = with_default_table(value, default_table)
511 2260 aaronmk
        return CompareCond(value)
512 2216 aaronmk
    else: return value
513 2219 aaronmk
514 2335 aaronmk
##### Joins
515
516 2352 aaronmk
join_same = object() # tells Join the left and right columns have the same name
517 2260 aaronmk
518 2353 aaronmk
# Tells Join the left and right columns have the same name and are never NULL
519
join_same_not_null = object()
520
521 2260 aaronmk
filter_out = object() # tells Join to filter out rows that match the join
522
523 2514 aaronmk
class Join(BasicObject):
524 2746 aaronmk
    def __init__(self, table, mapping={}, type_=None):
525 2260 aaronmk
        '''
526
        @param mapping dict(right_table_col=left_table_col, ...)
527 2352 aaronmk
            * if left_table_col is join_same: left_table_col = right_table_col
528 2353 aaronmk
              * Note that right_table_col must be a string
529
            * if left_table_col is join_same_not_null:
530
              left_table_col = right_table_col and both have NOT NULL constraint
531
              * Note that right_table_col must be a string
532 2260 aaronmk
        @param type_ None (for plain join)|str (e.g. 'LEFT')|filter_out
533
            * filter_out: equivalent to 'LEFT' with the query filtered by
534
              `table_pkey IS NULL` (indicating no match)
535
        '''
536
        if util.is_str(table): table = Table(table)
537
        assert type_ == None or util.is_str(type_) or type_ is filter_out
538
539
        self.table = table
540
        self.mapping = mapping
541
        self.type_ = type_
542
543 2749 aaronmk
    def to_str(self, db, left_table_):
544 2260 aaronmk
        def join(entry):
545
            '''Parses non-USING joins'''
546
            right_table_col, left_table_col = entry
547
548 2353 aaronmk
            # Switch order (right_table_col is on the left in the comparison)
549
            left = right_table_col
550
            right = left_table_col
551 2749 aaronmk
            left_table = self.table
552
            right_table = left_table_
553 2353 aaronmk
554 2747 aaronmk
            # Parse left side
555 2748 aaronmk
            left = with_default_table(left, left_table)
556 2747 aaronmk
557 2260 aaronmk
            # Parse special values
558 2747 aaronmk
            left_on_right = Col(left.name, right_table)
559
            if right is join_same: right = left_on_right
560 2353 aaronmk
            elif right is join_same_not_null:
561 2747 aaronmk
                right = CompareCond(left_on_right, '~=')
562 2260 aaronmk
563 2747 aaronmk
            # Parse right side
564 2353 aaronmk
            right = as_ValueCond(right, right_table)
565 2747 aaronmk
566
            return right.to_str(db, left)
567 2260 aaronmk
568 2265 aaronmk
        # Create join condition
569
        type_ = self.type_
570 2276 aaronmk
        joins = self.mapping
571 2746 aaronmk
        if joins == {}: join_cond = None
572
        elif type_ is not filter_out and reduce(operator.and_,
573 2460 aaronmk
            (v is join_same_not_null for v in joins.itervalues())):
574 2260 aaronmk
            # all cols w/ USING, so can use simpler USING syntax
575 2747 aaronmk
            cols = map(to_name_only_col, joins.iterkeys())
576
            join_cond = 'USING ('+(', '.join((c.to_str(db) for c in cols)))+')'
577 2757 aaronmk
        else: join_cond = combine_conds(map(join, joins.iteritems()), 'ON')
578 2260 aaronmk
579 2757 aaronmk
        if isinstance(self.table, NamedTable): whitespace = '\n'
580
        else: whitespace = ' '
581
582 2260 aaronmk
        # Create join
583
        if type_ is filter_out: type_ = 'LEFT'
584 2266 aaronmk
        str_ = ''
585
        if type_ != None: str_ += type_+' '
586 2757 aaronmk
        str_ += 'JOIN'+whitespace+self.table.to_str(db)
587
        if join_cond != None: str_ += whitespace+join_cond
588 2266 aaronmk
        return str_
589 2349 aaronmk
590 2514 aaronmk
    def __repr__(self): return self.to_str(mockDb, '<left_table>')
591 2424 aaronmk
592
##### Value exprs
593
594 3089 aaronmk
all_cols = CustomCode('*')
595
596 2737 aaronmk
default = CustomCode('DEFAULT')
597
598 3090 aaronmk
row_count = FunctionCall(InternalFunction('COUNT'), all_cols)
599 2674 aaronmk
600 3061 aaronmk
class Coalesce(FunctionCall):
601
    def __init__(self, *args):
602
        FunctionCall.__init__(self, InternalFunction('COALESCE'), *args)
603 3060 aaronmk
604 3062 aaronmk
class Nullif(FunctionCall):
605
    def __init__(self, *args):
606
        FunctionCall.__init__(self, InternalFunction('NULLIF'), *args)
607
608 2850 aaronmk
# See <http://www.postgresql.org/docs/8.3/static/datatype-numeric.html>
609 2958 aaronmk
null_sentinels = {
610
    'character varying': r'\N',
611
    'double precision': 'NaN',
612
    'integer': 2147483647,
613
    'text': r'\N',
614
    'timestamp with time zone': 'infinity'
615
}
616 2692 aaronmk
617 3061 aaronmk
class EnsureNotNull(Coalesce):
618 2850 aaronmk
    def __init__(self, value, type_):
619 3061 aaronmk
        Coalesce.__init__(self, as_Col(value),
620 2988 aaronmk
            Cast(type_, null_sentinels[type_]))
621 2850 aaronmk
622
        self.type = type_
623 3001 aaronmk
624
    def to_str(self, db):
625
        col = self.args[0]
626
        index_col_ = index_col(col)
627
        if index_col_ != None: return index_col_.to_str(db)
628 3061 aaronmk
        return Coalesce.to_str(self, db)
629 2850 aaronmk
630 2737 aaronmk
##### Table exprs
631
632
class Values(Code):
633
    def __init__(self, values):
634 2739 aaronmk
        '''
635
        @param values [...]|[[...], ...] Can be one or multiple rows.
636
        '''
637
        rows = values
638
        if len(values) >= 1 and not lists.is_seq(values[0]): # only one row
639
            rows = [values]
640
        for i, row in enumerate(rows):
641
            rows[i] = map(remove_col_rename, map(as_Value, row))
642 2737 aaronmk
643 2739 aaronmk
        self.rows = rows
644 2737 aaronmk
645
    def to_str(self, db):
646 2739 aaronmk
        def row_str(row):
647
            return '('+(', '.join((v.to_str(db) for v in row)))+')'
648
        return 'VALUES '+(', '.join(map(row_str, self.rows)))
649 2737 aaronmk
650 2740 aaronmk
def NamedValues(name, cols, values):
651 2745 aaronmk
    '''
652 3048 aaronmk
    @param cols None|[...]
653 2745 aaronmk
    @post `cols` will be changed to Col objects with the table set to `name`.
654
    '''
655 2834 aaronmk
    table = NamedTable(name, Values(values), cols)
656 3048 aaronmk
    if cols != None: set_cols_table(table, cols)
657 2834 aaronmk
    return table
658 2740 aaronmk
659 2674 aaronmk
##### Database structure
660
661
class TypedCol(Col):
662 2871 aaronmk
    def __init__(self, name, type_, default=None, nullable=True,
663
        constraints=None):
664 2818 aaronmk
        assert default == None or isinstance(default, Code)
665
666 2674 aaronmk
        Col.__init__(self, name)
667
668
        self.type = type_
669 2818 aaronmk
        self.default = default
670
        self.nullable = nullable
671 2871 aaronmk
        self.constraints = constraints
672 2674 aaronmk
673 2818 aaronmk
    def to_str(self, db):
674
        str_ = Col.to_str(self, db)+' '+self.type
675
        if not self.nullable: str_ += ' NOT NULL'
676
        if self.default != None: str_ += ' DEFAULT '+self.default.to_str(db)
677 2871 aaronmk
        if self.constraints != None: str_ += ' '+self.constraints
678 2818 aaronmk
        return str_
679 2674 aaronmk
680
    def to_Col(self): return Col(self.name)
681 2822 aaronmk
682 2840 aaronmk
ensure_not_null_excs = (NoUnderlyingTableException, KeyError)
683
684 2851 aaronmk
def ensure_not_null(db, col, type_=None):
685 2840 aaronmk
    '''
686 2855 aaronmk
    @param col If type_ is not set, must have an underlying column.
687 2851 aaronmk
    @param type_ If set, overrides the underlying column's type.
688 2840 aaronmk
    @return EnsureNotNull|Col
689
    @throws ensure_not_null_excs
690
    '''
691 2855 aaronmk
    nullable = True
692
    try: typed_col = db.col_info(underlying_col(col))
693
    except NoUnderlyingTableException:
694
        if type_ == None: raise
695
    else:
696
        if type_ == None: type_ = typed_col.type
697
        nullable = typed_col.nullable
698
699 2953 aaronmk
    if nullable:
700
        try: col = EnsureNotNull(col, type_)
701
        except KeyError, e:
702
            # Warn of no null sentinel for type, even if caller catches error
703
            warnings.warn(UserWarning(exc.str_(e)))
704
            raise
705
706 2840 aaronmk
    return col