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 3320 aaronmk
def unesc_name(name, quote='"'):
48
    removed_ref = [False]
49
    name = strings.remove_prefix(quote, name, removed_ref)
50
    if removed_ref[0]:
51
        name = strings.remove_suffix(quote, name, removed_ref)
52
        assert removed_ref[0]
53
        name = name.replace(quote+quote, quote)
54
    return name
55
56 2513 aaronmk
def clean_name(name): return name.replace('"', '').replace('`', '')
57
58 3041 aaronmk
def esc_comment(comment): return '/*'+comment.replace('*/', '* /')+'*/'
59
60 3182 aaronmk
def lstrip(str_):
61
    '''Also removes comments.'''
62
    if str_.startswith('/*'): comment, sep, str_ = str_.partition('*/')
63
    return str_.lstrip()
64
65 2659 aaronmk
##### General SQL code objects
66 2219 aaronmk
67 2349 aaronmk
class MockDb:
68 2503 aaronmk
    def esc_value(self, value): return strings.repr_no_u(value)
69 2349 aaronmk
70 2499 aaronmk
    def esc_name(self, name): return esc_name(name)
71 2859 aaronmk
72
    def col_info(self, col):
73
        return TypedCol(col.name, '<type>', CustomCode('<default>'), True)
74
75 2349 aaronmk
mockDb = MockDb()
76
77 2514 aaronmk
class BasicObject(objects.BasicObject):
78
    def __str__(self): return clean_name(strings.repr_no_u(self))
79
80 2659 aaronmk
##### Unparameterized code objects
81
82 2514 aaronmk
class Code(BasicObject):
83 3446 aaronmk
    def __init__(self, lang='sql'):
84
        self.lang = lang
85 3445 aaronmk
86 2658 aaronmk
    def to_str(self, db): raise NotImplementedError()
87 2349 aaronmk
88 2514 aaronmk
    def __repr__(self): return self.to_str(mockDb)
89 2211 aaronmk
90 2269 aaronmk
class CustomCode(Code):
91 3445 aaronmk
    def __init__(self, str_):
92
        Code.__init__(self)
93
94
        self.str_ = str_
95 2256 aaronmk
96
    def to_str(self, db): return self.str_
97
98 2815 aaronmk
def as_Code(value, db=None):
99
    '''
100
    @param db If set, runs db.std_code() on the value.
101
    '''
102 3447 aaronmk
    if isinstance(value, Code): return value
103
104 2815 aaronmk
    if util.is_str(value):
105
        if db != None: value = db.std_code(value)
106
        return CustomCode(value)
107 2659 aaronmk
    else: return Literal(value)
108
109 2540 aaronmk
class Expr(Code):
110 3445 aaronmk
    def __init__(self, expr):
111
        Code.__init__(self)
112
113
        self.expr = expr
114 2540 aaronmk
115
    def to_str(self, db): return '('+self.expr.to_str(db)+')'
116
117 3086 aaronmk
##### Names
118
119
class Name(Code):
120 3087 aaronmk
    def __init__(self, name):
121 3445 aaronmk
        Code.__init__(self)
122
123 3087 aaronmk
        name = truncate(name)
124
125
        self.name = name
126 3086 aaronmk
127
    def to_str(self, db): return db.esc_name(self.name)
128
129
def as_Name(value):
130
    if isinstance(value, Code): return value
131
    else: return Name(value)
132
133 2335 aaronmk
##### Literal values
134
135 2216 aaronmk
class Literal(Code):
136 3445 aaronmk
    def __init__(self, value):
137
        Code.__init__(self)
138
139
        self.value = value
140 2213 aaronmk
141
    def to_str(self, db): return db.esc_value(self.value)
142 2211 aaronmk
143 2400 aaronmk
def as_Value(value):
144
    if isinstance(value, Code): return value
145
    else: return Literal(value)
146
147 3429 aaronmk
def is_literal(value): return isinstance(value, Literal)
148 2216 aaronmk
149 3429 aaronmk
def is_null(value): return is_literal(value) and value.value == None
150
151 2711 aaronmk
##### Derived elements
152
153
src_self = object() # tells Col that it is its own source column
154
155
class Derived(Code):
156
    def __init__(self, srcs):
157 2712 aaronmk
        '''An element which was derived from some other element(s).
158 2711 aaronmk
        @param srcs See self.set_srcs()
159
        '''
160 3445 aaronmk
        Code.__init__(self)
161
162 2711 aaronmk
        self.set_srcs(srcs)
163
164 2713 aaronmk
    def set_srcs(self, srcs, overwrite=True):
165 2711 aaronmk
        '''
166
        @param srcs (self_type...)|src_self The element(s) this is derived from
167
        '''
168 2713 aaronmk
        if not overwrite and self.srcs != (): return # already set
169
170 2711 aaronmk
        if srcs == src_self: srcs = (self,)
171
        srcs = tuple(srcs) # make Col hashable
172
        self.srcs = srcs
173
174
    def _compare_on(self):
175
        compare_on = self.__dict__.copy()
176
        del compare_on['srcs'] # ignore
177
        return compare_on
178
179
def cols_srcs(cols): return lists.uniqify(iters.flatten((v.srcs for v in cols)))
180
181 2335 aaronmk
##### Tables
182
183 2712 aaronmk
class Table(Derived):
184 2991 aaronmk
    def __init__(self, name, schema=None, srcs=(), is_temp=False):
185 2211 aaronmk
        '''
186
        @param schema str|None (for no schema)
187 2712 aaronmk
        @param srcs (Table...)|src_self See Derived.set_srcs()
188 2211 aaronmk
        '''
189 2712 aaronmk
        Derived.__init__(self, srcs)
190
191 3091 aaronmk
        if util.is_str(name): name = truncate(name)
192 2843 aaronmk
193 2211 aaronmk
        self.name = name
194
        self.schema = schema
195 2991 aaronmk
        self.is_temp = is_temp
196 3000 aaronmk
        self.index_cols = {}
197 2211 aaronmk
198 2348 aaronmk
    def to_str(self, db):
199
        str_ = ''
200 3088 aaronmk
        if self.schema != None: str_ += as_Name(self.schema).to_str(db)+'.'
201
        str_ += as_Name(self.name).to_str(db)
202 2348 aaronmk
        return str_
203 2336 aaronmk
204
    def to_Table(self): return self
205 3000 aaronmk
206
    def _compare_on(self):
207
        compare_on = Derived._compare_on(self)
208
        del compare_on['index_cols'] # ignore
209
        return compare_on
210 2211 aaronmk
211 2835 aaronmk
def is_underlying_table(table):
212
    return isinstance(table, Table) and table.to_Table() is table
213 2832 aaronmk
214 2902 aaronmk
class NoUnderlyingTableException(Exception): pass
215
216
def underlying_table(table):
217
    table = remove_table_rename(table)
218
    if not is_underlying_table(table): raise NoUnderlyingTableException
219
    return table
220
221 2776 aaronmk
def as_Table(table, schema=None):
222 2270 aaronmk
    if table == None or isinstance(table, Code): return table
223 2776 aaronmk
    else: return Table(table, schema)
224 2219 aaronmk
225 3101 aaronmk
def suffixed_table(table, suffix):
226 3128 aaronmk
    table = copy.copy(table) # don't modify input!
227
    table.name = concat(table.name, suffix)
228
    return table
229 2707 aaronmk
230 2336 aaronmk
class NamedTable(Table):
231
    def __init__(self, name, code, cols=None):
232
        Table.__init__(self, name)
233
234 3016 aaronmk
        code = as_Table(code)
235 2741 aaronmk
        if not isinstance(code, (Table, FunctionCall, Expr)): code = Expr(code)
236 3020 aaronmk
        if cols != None: cols = [to_name_only_col(c).to_Col() for c in cols]
237 2336 aaronmk
238
        self.code = code
239
        self.cols = cols
240
241
    def to_str(self, db):
242 3026 aaronmk
        str_ = self.code.to_str(db)
243
        if str_.find('\n') >= 0: whitespace = '\n'
244
        else: whitespace = ' '
245
        str_ += whitespace+'AS '+Table.to_str(self, db)
246 2742 aaronmk
        if self.cols != None:
247
            str_ += ' ('+(', '.join((c.to_str(db) for c in self.cols)))+')'
248 2336 aaronmk
        return str_
249
250
    def to_Table(self): return Table(self.name)
251
252 2753 aaronmk
def remove_table_rename(table):
253
    if isinstance(table, NamedTable): table = table.code
254
    return table
255
256 2335 aaronmk
##### Columns
257
258 2711 aaronmk
class Col(Derived):
259 2701 aaronmk
    def __init__(self, name, table=None, srcs=()):
260 2211 aaronmk
        '''
261
        @param table Table|None (for no table)
262 2711 aaronmk
        @param srcs (Col...)|src_self See Derived.set_srcs()
263 2211 aaronmk
        '''
264 2711 aaronmk
        Derived.__init__(self, srcs)
265
266 3091 aaronmk
        if util.is_str(name): name = truncate(name)
267 2241 aaronmk
        if util.is_str(table): table = Table(table)
268 2211 aaronmk
        assert table == None or isinstance(table, Table)
269
270
        self.name = name
271
        self.table = table
272
273 2989 aaronmk
    def to_str(self, db, for_str=False):
274 3088 aaronmk
        str_ = as_Name(self.name).to_str(db)
275 2989 aaronmk
        if for_str: str_ = clean_name(str_)
276 2933 aaronmk
        if self.table != None:
277 2989 aaronmk
            table = self.table.to_Table()
278
            if for_str: str_ = concat(str(table), '.'+str_)
279
            else: str_ = table.to_str(db)+'.'+str_
280 2211 aaronmk
        return str_
281 2314 aaronmk
282 2989 aaronmk
    def __str__(self): return self.to_str(mockDb, for_str=True)
283 2933 aaronmk
284 2314 aaronmk
    def to_Col(self): return self
285 2211 aaronmk
286 2767 aaronmk
def is_table_col(col): return isinstance(col, Col) and col.table != None
287 2393 aaronmk
288 3000 aaronmk
def index_col(col):
289
    if not is_table_col(col): return None
290 3104 aaronmk
291
    table = col.table
292
    try: name = table.index_cols[col.name]
293
    except KeyError: return None
294
    else: return Col(name, table, col.srcs)
295 2999 aaronmk
296 3024 aaronmk
def is_temp_col(col): return is_table_col(col) and col.table.is_temp
297 2996 aaronmk
298 2563 aaronmk
def as_Col(col, table=None, name=None):
299
    '''
300
    @param name If not None, any non-Col input will be renamed using NamedCol.
301
    '''
302
    if name != None:
303
        col = as_Value(col)
304
        if not isinstance(col, Col): col = NamedCol(name, col)
305 2333 aaronmk
306
    if isinstance(col, Code): return col
307 2260 aaronmk
    else: return Col(col, table)
308
309 3093 aaronmk
def with_table(col, table):
310 3105 aaronmk
    if isinstance(col, NamedCol): pass # doesn't take a table
311
    elif isinstance(col, FunctionCall):
312
        col = copy.deepcopy(col) # don't modify input!
313
        col.args[0].table = table
314
    else:
315 3098 aaronmk
        col = copy.copy(col) # don't modify input!
316
        col.table = table
317 3093 aaronmk
    return col
318
319 3100 aaronmk
def with_default_table(col, table):
320 2747 aaronmk
    col = as_Col(col)
321 3100 aaronmk
    if col.table == None: col = with_table(col, table)
322 2747 aaronmk
    return col
323
324 2744 aaronmk
def set_cols_table(table, cols):
325
    table = as_Table(table)
326
327
    for i, col in enumerate(cols):
328
        col = cols[i] = as_Col(col)
329
        col.table = table
330
331 2401 aaronmk
def to_name_only_col(col, check_table=None):
332
    col = as_Col(col)
333 3020 aaronmk
    if not is_table_col(col): return col
334 2401 aaronmk
335
    if check_table != None:
336
        table = col.table
337
        assert table == None or table == check_table
338
    return Col(col.name)
339
340 2993 aaronmk
def suffixed_col(col, suffix):
341
    return Col(concat(col.name, suffix), col.table, col.srcs)
342
343 2323 aaronmk
class NamedCol(Col):
344 2229 aaronmk
    def __init__(self, name, code):
345 2310 aaronmk
        Col.__init__(self, name)
346
347 3016 aaronmk
        code = as_Value(code)
348 2229 aaronmk
349
        self.code = code
350
351
    def to_str(self, db):
352 2310 aaronmk
        return self.code.to_str(db)+' AS '+Col.to_str(self, db)
353 2314 aaronmk
354
    def to_Col(self): return Col(self.name)
355 2229 aaronmk
356 2462 aaronmk
def remove_col_rename(col):
357
    if isinstance(col, NamedCol): col = col.code
358
    return col
359
360 2830 aaronmk
def underlying_col(col):
361
    col = remove_col_rename(col)
362 2849 aaronmk
    if not isinstance(col, Col): raise NoUnderlyingTableException
363
364 2902 aaronmk
    return Col(col.name, underlying_table(col.table), col.srcs)
365 2830 aaronmk
366 2703 aaronmk
def wrap(wrap_func, value):
367
    '''Wraps a value, propagating any column renaming to the returned value.'''
368
    if isinstance(value, NamedCol):
369
        return NamedCol(value.name, wrap_func(value.code))
370
    else: return wrap_func(value)
371
372 2667 aaronmk
class ColDict(dicts.DictProxy):
373 2564 aaronmk
    '''A dict that automatically makes inserted entries Col objects'''
374
375 2645 aaronmk
    def __init__(self, db, keys_table, dict_={}):
376 3158 aaronmk
        dicts.DictProxy.__init__(self, OrderedDict())
377 2667 aaronmk
378 2645 aaronmk
        keys_table = as_Table(keys_table)
379
380 2642 aaronmk
        self.db = db
381 2641 aaronmk
        self.table = keys_table
382 2653 aaronmk
        self.update(dict_) # after setting vars because __setitem__() needs them
383 2641 aaronmk
384 2667 aaronmk
    def copy(self): return ColDict(self.db, self.table, self.inner.copy())
385 2655 aaronmk
386 2667 aaronmk
    def __getitem__(self, key):
387
        return dicts.DictProxy.__getitem__(self, self._key(key))
388 2653 aaronmk
389 2564 aaronmk
    def __setitem__(self, key, value):
390 2642 aaronmk
        key = self._key(key)
391 2819 aaronmk
        if value == None: value = self.db.col_info(key).default
392 2667 aaronmk
        dicts.DictProxy.__setitem__(self, key, as_Col(value, name=key.name))
393 2564 aaronmk
394 2641 aaronmk
    def _key(self, key): return as_Col(key, self.table)
395 2564 aaronmk
396 2524 aaronmk
##### Functions
397
398 2912 aaronmk
Function = Table
399 2911 aaronmk
as_Function = as_Table
400
401 2691 aaronmk
class InternalFunction(CustomCode): pass
402
403 3442 aaronmk
#### Calls
404
405 2941 aaronmk
class NamedArg(NamedCol):
406
    def __init__(self, name, value):
407
        NamedCol.__init__(self, name, value)
408
409
    def to_str(self, db):
410
        return Col.to_str(self, db)+' := '+self.code.to_str(db)
411
412 2524 aaronmk
class FunctionCall(Code):
413 2941 aaronmk
    def __init__(self, function, *args, **kw_args):
414 2524 aaronmk
        '''
415 2690 aaronmk
        @param args [Code|literal-value...] The function's arguments
416 2524 aaronmk
        '''
417 3445 aaronmk
        Code.__init__(self)
418
419 3016 aaronmk
        function = as_Function(function)
420 2941 aaronmk
        def filter_(arg): return remove_col_rename(as_Value(arg))
421
        args = map(filter_, args)
422
        args += [NamedArg(k, filter_(v)) for k, v in kw_args.iteritems()]
423 2524 aaronmk
424
        self.function = function
425
        self.args = args
426
427
    def to_str(self, db):
428
        args_str = ', '.join((v.to_str(db) for v in self.args))
429
        return self.function.to_str(db)+'('+args_str+')'
430
431 2533 aaronmk
def wrap_in_func(function, value):
432
    '''Wraps a value inside a function call.
433
    Propagates any column renaming to the returned value.
434
    '''
435 2703 aaronmk
    return wrap(lambda v: FunctionCall(function, v), value)
436 2533 aaronmk
437 2561 aaronmk
def unwrap_func_call(func_call, check_name=None):
438
    '''Unwraps any function call to its first argument.
439
    Also removes any column renaming.
440
    '''
441
    func_call = remove_col_rename(func_call)
442
    if not isinstance(func_call, FunctionCall): return func_call
443
444
    if check_name != None:
445
        name = func_call.function.name
446
        assert name == None or name == check_name
447
    return func_call.args[0]
448
449 3442 aaronmk
#### Definitions
450
451
class FunctionDef(Code):
452 3456 aaronmk
    def __init__(self, function, return_type, body, modifiers=None):
453 3445 aaronmk
        Code.__init__(self)
454
455 3444 aaronmk
        body = as_Code(body)
456
457 3442 aaronmk
        self.function = function
458
        self.return_type = return_type
459
        self.body = body
460 3456 aaronmk
        self.modifiers = modifiers
461 3442 aaronmk
462
    def to_str(self, db):
463
        str_ = '''\
464
CREATE FUNCTION '''+self.function.to_str(db)+'''()
465 3451 aaronmk
RETURNS '''+self.return_type+'''
466 3448 aaronmk
LANGUAGE '''+self.body.lang+'''
467 3456 aaronmk
'''
468
        if self.modifiers != None: str_ += self.modifiers+'\n'
469
        str_ += '''\
470 3442 aaronmk
AS $$
471 3444 aaronmk
'''+self.body.to_str(db)+'''
472 3442 aaronmk
$$;
473
'''
474
        return str_
475
476 3454 aaronmk
### PL/pgSQL
477
478
class ExcHandler(BasicObject):
479
    def __init__(self, exc, handler=None):
480
        if handler != None: handler = as_Code(handler)
481
482
        self.exc = exc
483
        self.handler = handler
484
485
    def to_str(self, db, body):
486
        body = as_Code(body)
487
488
        if self.handler != None: handler_str = '\n'+self.handler.to_str(db)
489
        else: handler_str = ' NULL;'
490
491
        str_ = '''\
492
BEGIN
493 3458 aaronmk
'''+body.to_str(db)+'''\
494 3454 aaronmk
EXCEPTION
495
    WHEN '''+self.exc+''' THEN'''+handler_str+'''
496
END;\
497
'''
498
        return str_
499 3457 aaronmk
500
    def __repr__(self): return self.to_str(mockDb, '<body>')
501 3454 aaronmk
502
unique_violation_handler = ExcHandler('unique_violation')
503
504 3449 aaronmk
class RowExcIgnore(Code):
505
    def __init__(self, row_type, select_query, with_row, cols=None,
506 3455 aaronmk
        exc_handler=unique_violation_handler, row_var='row'):
507 3449 aaronmk
        Code.__init__(self, lang='plpgsql')
508
509
        select_query = as_Code(select_query)
510
        with_row = as_Code(with_row)
511 3452 aaronmk
        row_var = as_Table(row_var)
512 3449 aaronmk
513
        self.row_type = row_type
514
        self.select_query = select_query
515
        self.with_row = with_row
516
        self.cols = cols
517 3455 aaronmk
        self.exc_handler = exc_handler
518 3452 aaronmk
        self.row_var = row_var
519 3449 aaronmk
520
    def to_str(self, db):
521 3452 aaronmk
        if self.cols == None: row_vars = [self.row_var]
522
        else: row_vars = [Col(c.name, self.row_var) for c in self.cols]
523 3449 aaronmk
524
        str_ = '''\
525
DECLARE
526 3452 aaronmk
    '''+self.row_var.to_str(db)+''' '''+self.row_type+''';
527 3449 aaronmk
BEGIN
528
    /* Need an EXCEPTION block for each individual row because "When an error is
529
    caught by an EXCEPTION clause, [...] all changes to persistent database
530
    state within the block are rolled back."
531
    This is unfortunate because "A block containing an EXCEPTION clause is
532
    significantly more expensive to enter and exit than a block without one."
533
    (http://www.postgresql.org/docs/8.3/static/plpgsql-control-structures.html\
534
#PLPGSQL-ERROR-TRAPPING)
535
    */
536
    FOR '''+(', '.join((v.to_str(db) for v in row_vars)))+''' IN
537
'''+self.select_query.to_str(db)+'''
538
    LOOP
539 3455 aaronmk
'''+self.exc_handler.to_str(db, self.with_row)+'''
540 3449 aaronmk
    END LOOP;
541
END;\
542
'''
543
        return str_
544
545 2986 aaronmk
##### Casts
546
547
class Cast(FunctionCall):
548
    def __init__(self, type_, value):
549
        value = as_Value(value)
550
551
        self.type_ = type_
552
        self.value = value
553
554
    def to_str(self, db):
555
        return 'CAST('+self.value.to_str(db)+' AS '+self.type_+')'
556
557 3354 aaronmk
def cast_literal(value):
558 3429 aaronmk
    if not is_literal(value): return value
559 3354 aaronmk
560
    if util.is_str(value.value): value = Cast('text', value)
561
    return value
562
563 2335 aaronmk
##### Conditions
564 2259 aaronmk
565 3350 aaronmk
class NotCond(Code):
566
    def __init__(self, cond):
567 3445 aaronmk
        Code.__init__(self)
568
569 3350 aaronmk
        self.cond = cond
570
571
    def to_str(self, db): return 'NOT '+self.cond.to_str(db)
572
573 2398 aaronmk
class ColValueCond(Code):
574
    def __init__(self, col, value):
575 3445 aaronmk
        Code.__init__(self)
576
577 2398 aaronmk
        value = as_ValueCond(value)
578
579
        self.col = col
580
        self.value = value
581
582
    def to_str(self, db): return self.value.to_str(db, self.col)
583
584 2577 aaronmk
def combine_conds(conds, keyword=None):
585
    '''
586
    @param keyword The keyword to add before the conditions, if any
587
    '''
588
    str_ = ''
589
    if keyword != None:
590
        if conds == []: whitespace = ''
591
        elif len(conds) == 1: whitespace = ' '
592
        else: whitespace = '\n'
593
        str_ += keyword+whitespace
594
595
    str_ += '\nAND '.join(conds)
596
    return str_
597
598 2398 aaronmk
##### Condition column comparisons
599
600 2514 aaronmk
class ValueCond(BasicObject):
601 2213 aaronmk
    def __init__(self, value):
602 2858 aaronmk
        value = remove_col_rename(as_Value(value))
603 2213 aaronmk
604
        self.value = value
605 2214 aaronmk
606 2216 aaronmk
    def to_str(self, db, left_value):
607 2214 aaronmk
        '''
608 2216 aaronmk
        @param left_value The Code object that the condition is being applied on
609 2214 aaronmk
        '''
610
        raise NotImplemented()
611 2228 aaronmk
612 2514 aaronmk
    def __repr__(self): return self.to_str(mockDb, '<left_value>')
613 2211 aaronmk
614
class CompareCond(ValueCond):
615
    def __init__(self, value, operator='='):
616 2222 aaronmk
        '''
617
        @param operator By default, compares NULL values literally. Use '~=' or
618
            '~!=' to pass NULLs through.
619
        '''
620 2211 aaronmk
        ValueCond.__init__(self, value)
621
        self.operator = operator
622
623 2216 aaronmk
    def to_str(self, db, left_value):
624 2858 aaronmk
        left_value = remove_col_rename(as_Col(left_value))
625 2216 aaronmk
626 2222 aaronmk
        right_value = self.value
627
628
        # Parse operator
629 2216 aaronmk
        operator = self.operator
630 2222 aaronmk
        passthru_null_ref = [False]
631
        operator = strings.remove_prefix('~', operator, passthru_null_ref)
632
        neg_ref = [False]
633
        operator = strings.remove_prefix('!', operator, neg_ref)
634 2844 aaronmk
        equals = operator.endswith('=') # also includes <=, >=
635 2222 aaronmk
636 2825 aaronmk
        # Handle nullable columns
637
        check_null = False
638 2844 aaronmk
        if not passthru_null_ref[0]: # NULLs compare equal
639 2857 aaronmk
            try: left_value = ensure_not_null(db, left_value)
640 2844 aaronmk
            except ensure_not_null_excs: # fall back to alternate method
641
                check_null = equals and isinstance(right_value, Col)
642 2837 aaronmk
            else:
643 2857 aaronmk
                if isinstance(left_value, EnsureNotNull):
644
                    right_value = ensure_not_null(db, right_value,
645
                        left_value.type) # apply same function to both sides
646 2825 aaronmk
647 2844 aaronmk
        if equals and is_null(right_value): operator = 'IS'
648
649 2825 aaronmk
        left = left_value.to_str(db)
650
        right = right_value.to_str(db)
651
652 2222 aaronmk
        # Create str
653
        str_ = left+' '+operator+' '+right
654 2825 aaronmk
        if check_null:
655 2578 aaronmk
            str_ = '('+str_+' OR ('+left+' IS NULL AND '+right+' IS NULL))'
656
        if neg_ref[0]: str_ = 'NOT '+str_
657 2222 aaronmk
        return str_
658 2216 aaronmk
659 2260 aaronmk
# Tells as_ValueCond() to assume a non-ValueCond is a literal value
660
assume_literal = object()
661
662
def as_ValueCond(value, default_table=assume_literal):
663
    if not isinstance(value, ValueCond):
664
        if default_table is not assume_literal:
665 2748 aaronmk
            value = with_default_table(value, default_table)
666 2260 aaronmk
        return CompareCond(value)
667 2216 aaronmk
    else: return value
668 2219 aaronmk
669 2335 aaronmk
##### Joins
670
671 2352 aaronmk
join_same = object() # tells Join the left and right columns have the same name
672 2260 aaronmk
673 2353 aaronmk
# Tells Join the left and right columns have the same name and are never NULL
674
join_same_not_null = object()
675
676 2260 aaronmk
filter_out = object() # tells Join to filter out rows that match the join
677
678 2514 aaronmk
class Join(BasicObject):
679 2746 aaronmk
    def __init__(self, table, mapping={}, type_=None):
680 2260 aaronmk
        '''
681
        @param mapping dict(right_table_col=left_table_col, ...)
682 2352 aaronmk
            * if left_table_col is join_same: left_table_col = right_table_col
683 2353 aaronmk
              * Note that right_table_col must be a string
684
            * if left_table_col is join_same_not_null:
685
              left_table_col = right_table_col and both have NOT NULL constraint
686
              * Note that right_table_col must be a string
687 2260 aaronmk
        @param type_ None (for plain join)|str (e.g. 'LEFT')|filter_out
688
            * filter_out: equivalent to 'LEFT' with the query filtered by
689
              `table_pkey IS NULL` (indicating no match)
690
        '''
691
        if util.is_str(table): table = Table(table)
692
        assert type_ == None or util.is_str(type_) or type_ is filter_out
693
694
        self.table = table
695
        self.mapping = mapping
696
        self.type_ = type_
697
698 2749 aaronmk
    def to_str(self, db, left_table_):
699 2260 aaronmk
        def join(entry):
700
            '''Parses non-USING joins'''
701
            right_table_col, left_table_col = entry
702
703 2353 aaronmk
            # Switch order (right_table_col is on the left in the comparison)
704
            left = right_table_col
705
            right = left_table_col
706 2749 aaronmk
            left_table = self.table
707
            right_table = left_table_
708 2353 aaronmk
709 2747 aaronmk
            # Parse left side
710 2748 aaronmk
            left = with_default_table(left, left_table)
711 2747 aaronmk
712 2260 aaronmk
            # Parse special values
713 2747 aaronmk
            left_on_right = Col(left.name, right_table)
714
            if right is join_same: right = left_on_right
715 2353 aaronmk
            elif right is join_same_not_null:
716 2747 aaronmk
                right = CompareCond(left_on_right, '~=')
717 2260 aaronmk
718 2747 aaronmk
            # Parse right side
719 2353 aaronmk
            right = as_ValueCond(right, right_table)
720 2747 aaronmk
721
            return right.to_str(db, left)
722 2260 aaronmk
723 2265 aaronmk
        # Create join condition
724
        type_ = self.type_
725 2276 aaronmk
        joins = self.mapping
726 2746 aaronmk
        if joins == {}: join_cond = None
727
        elif type_ is not filter_out and reduce(operator.and_,
728 2460 aaronmk
            (v is join_same_not_null for v in joins.itervalues())):
729 2260 aaronmk
            # all cols w/ USING, so can use simpler USING syntax
730 2747 aaronmk
            cols = map(to_name_only_col, joins.iterkeys())
731
            join_cond = 'USING ('+(', '.join((c.to_str(db) for c in cols)))+')'
732 2757 aaronmk
        else: join_cond = combine_conds(map(join, joins.iteritems()), 'ON')
733 2260 aaronmk
734 2757 aaronmk
        if isinstance(self.table, NamedTable): whitespace = '\n'
735
        else: whitespace = ' '
736
737 2260 aaronmk
        # Create join
738
        if type_ is filter_out: type_ = 'LEFT'
739 2266 aaronmk
        str_ = ''
740
        if type_ != None: str_ += type_+' '
741 2757 aaronmk
        str_ += 'JOIN'+whitespace+self.table.to_str(db)
742
        if join_cond != None: str_ += whitespace+join_cond
743 2266 aaronmk
        return str_
744 2349 aaronmk
745 2514 aaronmk
    def __repr__(self): return self.to_str(mockDb, '<left_table>')
746 2424 aaronmk
747
##### Value exprs
748
749 3089 aaronmk
all_cols = CustomCode('*')
750
751 2737 aaronmk
default = CustomCode('DEFAULT')
752
753 3090 aaronmk
row_count = FunctionCall(InternalFunction('COUNT'), all_cols)
754 2674 aaronmk
755 3061 aaronmk
class Coalesce(FunctionCall):
756
    def __init__(self, *args):
757
        FunctionCall.__init__(self, InternalFunction('COALESCE'), *args)
758 3060 aaronmk
759 3062 aaronmk
class Nullif(FunctionCall):
760
    def __init__(self, *args):
761
        FunctionCall.__init__(self, InternalFunction('NULLIF'), *args)
762
763 2850 aaronmk
# See <http://www.postgresql.org/docs/8.3/static/datatype-numeric.html>
764 2958 aaronmk
null_sentinels = {
765
    'character varying': r'\N',
766
    'double precision': 'NaN',
767
    'integer': 2147483647,
768
    'text': r'\N',
769
    'timestamp with time zone': 'infinity'
770
}
771 2692 aaronmk
772 3061 aaronmk
class EnsureNotNull(Coalesce):
773 2850 aaronmk
    def __init__(self, value, type_):
774 3061 aaronmk
        Coalesce.__init__(self, as_Col(value),
775 2988 aaronmk
            Cast(type_, null_sentinels[type_]))
776 2850 aaronmk
777
        self.type = type_
778 3001 aaronmk
779
    def to_str(self, db):
780
        col = self.args[0]
781
        index_col_ = index_col(col)
782
        if index_col_ != None: return index_col_.to_str(db)
783 3061 aaronmk
        return Coalesce.to_str(self, db)
784 2850 aaronmk
785 2737 aaronmk
##### Table exprs
786
787
class Values(Code):
788
    def __init__(self, values):
789 2739 aaronmk
        '''
790
        @param values [...]|[[...], ...] Can be one or multiple rows.
791
        '''
792 3445 aaronmk
        Code.__init__(self)
793
794 2739 aaronmk
        rows = values
795
        if len(values) >= 1 and not lists.is_seq(values[0]): # only one row
796
            rows = [values]
797
        for i, row in enumerate(rows):
798
            rows[i] = map(remove_col_rename, map(as_Value, row))
799 2737 aaronmk
800 2739 aaronmk
        self.rows = rows
801 2737 aaronmk
802
    def to_str(self, db):
803 2739 aaronmk
        def row_str(row):
804
            return '('+(', '.join((v.to_str(db) for v in row)))+')'
805
        return 'VALUES '+(', '.join(map(row_str, self.rows)))
806 2737 aaronmk
807 2740 aaronmk
def NamedValues(name, cols, values):
808 2745 aaronmk
    '''
809 3048 aaronmk
    @param cols None|[...]
810 2745 aaronmk
    @post `cols` will be changed to Col objects with the table set to `name`.
811
    '''
812 2834 aaronmk
    table = NamedTable(name, Values(values), cols)
813 3048 aaronmk
    if cols != None: set_cols_table(table, cols)
814 2834 aaronmk
    return table
815 2740 aaronmk
816 2674 aaronmk
##### Database structure
817
818
class TypedCol(Col):
819 2871 aaronmk
    def __init__(self, name, type_, default=None, nullable=True,
820
        constraints=None):
821 2818 aaronmk
        assert default == None or isinstance(default, Code)
822
823 2674 aaronmk
        Col.__init__(self, name)
824
825
        self.type = type_
826 2818 aaronmk
        self.default = default
827
        self.nullable = nullable
828 2871 aaronmk
        self.constraints = constraints
829 2674 aaronmk
830 2818 aaronmk
    def to_str(self, db):
831
        str_ = Col.to_str(self, db)+' '+self.type
832
        if not self.nullable: str_ += ' NOT NULL'
833
        if self.default != None: str_ += ' DEFAULT '+self.default.to_str(db)
834 2871 aaronmk
        if self.constraints != None: str_ += ' '+self.constraints
835 2818 aaronmk
        return str_
836 2674 aaronmk
837
    def to_Col(self): return Col(self.name)
838 2822 aaronmk
839 2840 aaronmk
ensure_not_null_excs = (NoUnderlyingTableException, KeyError)
840
841 2851 aaronmk
def ensure_not_null(db, col, type_=None):
842 2840 aaronmk
    '''
843 2855 aaronmk
    @param col If type_ is not set, must have an underlying column.
844 2851 aaronmk
    @param type_ If set, overrides the underlying column's type.
845 2840 aaronmk
    @return EnsureNotNull|Col
846
    @throws ensure_not_null_excs
847
    '''
848 2855 aaronmk
    nullable = True
849
    try: typed_col = db.col_info(underlying_col(col))
850
    except NoUnderlyingTableException:
851 3355 aaronmk
        col = remove_col_rename(col)
852 3429 aaronmk
        if is_literal(col) and not is_null(col): nullable = False
853 3355 aaronmk
        elif type_ == None: raise
854 2855 aaronmk
    else:
855
        if type_ == None: type_ = typed_col.type
856
        nullable = typed_col.nullable
857
858 2953 aaronmk
    if nullable:
859
        try: col = EnsureNotNull(col, type_)
860
        except KeyError, e:
861
            # Warn of no null sentinel for type, even if caller catches error
862
            warnings.warn(UserWarning(exc.str_(e)))
863
            raise
864
865 2840 aaronmk
    return col