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