Project

General

Profile

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