Project

General

Profile

1
# SQL code generation
2

    
3
import copy
4
import itertools
5
import operator
6
from ordereddict import OrderedDict
7
import re
8
import UserDict
9
import warnings
10

    
11
import dicts
12
import exc
13
import iters
14
import lists
15
import objects
16
import regexp
17
import strings
18
import util
19

    
20
##### Names
21

    
22
identifier_max_len = 63 # works for both PostgreSQL and MySQL
23

    
24
def concat(str_, suffix):
25
    '''Preserves version so that it won't be truncated off the string, leading
26
    to collisions.'''
27
    # Preserve version
28
    match = re.match(r'^(.*?)((?:(?:#\d+)?\)?)*(?:\.\w+)+(?:::[\w ]+)*)$', str_)
29
    if match:
30
        str_, old_suffix = match.groups()
31
        suffix = old_suffix+suffix
32
    
33
    return strings.concat(str_, suffix, identifier_max_len)
34

    
35
def truncate(str_): return concat(str_, '')
36

    
37
def is_safe_name(name):
38
    '''A name is safe *and unambiguous* if it:
39
    * contains only *lowercase* word (\w) characters
40
    * doesn't start with a digit
41
    * contains "_", so that it's not a keyword
42
    '''
43
    return re.match(r'^(?=.*_)(?!\d)[^\WA-Z]+$', name)
44

    
45
def esc_name(name, quote='"'):
46
    return quote + name.replace(quote, quote+quote) + quote
47
        # doubling an embedded quote escapes it in both PostgreSQL and MySQL
48

    
49
def unesc_name(name, quote='"'):
50
    removed_ref = [False]
51
    name = strings.remove_prefix(quote, name, removed_ref)
52
    if removed_ref[0]:
53
        name = strings.remove_suffix(quote, name, removed_ref)
54
        assert removed_ref[0]
55
        name = name.replace(quote+quote, quote)
56
    return name
57

    
58
def clean_name(name): return name.replace('"', '').replace('`', '')
59

    
60
def esc_comment(comment): return '/*'+comment.replace('*/', '* /')+'*/'
61

    
62
def lstrip(str_):
63
    '''Also removes comments.'''
64
    if str_.startswith('/*'): comment, sep, str_ = str_.partition('*/')
65
    return str_.lstrip()
66

    
67
##### General SQL code objects
68

    
69
class MockDb:
70
    def esc_value(self, value): return strings.repr_no_u(value)
71
    
72
    def esc_name(self, name): return esc_name(name)
73
    
74
    def col_info(self, col):
75
        return TypedCol(col.name, '<type>', CustomCode('<default>'), True)
76

    
77
mockDb = MockDb()
78

    
79
class BasicObject(objects.BasicObject):
80
    def __str__(self): return clean_name(strings.repr_no_u(self))
81

    
82
##### Unparameterized code objects
83

    
84
class Code(BasicObject):
85
    def __init__(self, lang='sql'):
86
        self.lang = lang
87
    
88
    def to_str(self, db): raise NotImplementedError()
89
    
90
    def __repr__(self): return self.to_str(mockDb)
91

    
92
class CustomCode(Code):
93
    def __init__(self, str_):
94
        Code.__init__(self)
95
        
96
        self.str_ = str_
97
    
98
    def to_str(self, db): return self.str_
99

    
100
def as_Code(value, db=None):
101
    '''
102
    @param db If set, runs db.std_code() on the value.
103
    '''
104
    if isinstance(value, Code): return value
105
    
106
    if util.is_str(value):
107
        if db != None: value = db.std_code(value)
108
        return CustomCode(value)
109
    else: return Literal(value)
110

    
111
class Expr(Code):
112
    def __init__(self, expr):
113
        Code.__init__(self)
114
        
115
        self.expr = expr
116
    
117
    def to_str(self, db): return '('+self.expr.to_str(db)+')'
118

    
119
##### Names
120

    
121
class Name(Code):
122
    def __init__(self, name):
123
        Code.__init__(self)
124
        
125
        name = truncate(name)
126
        
127
        self.name = name
128
    
129
    def to_str(self, db): return db.esc_name(self.name)
130

    
131
def as_Name(value):
132
    if isinstance(value, Code): return value
133
    else: return Name(value)
134

    
135
##### Literal values
136

    
137
#### Primitives
138

    
139
class Literal(Code):
140
    def __init__(self, value):
141
        Code.__init__(self)
142
        
143
        self.value = value
144
    
145
    def to_str(self, db): return db.esc_value(self.value)
146

    
147
def as_Value(value):
148
    if isinstance(value, Code): return value
149
    else: return Literal(value)
150

    
151
def get_value(value):
152
    '''Unwraps a Literal's value'''
153
    value = remove_col_rename(value)
154
    if isinstance(value, Literal): return value.value
155
    else:
156
        assert not isinstance(value, Code)
157
        return value
158

    
159
def is_literal(value): return isinstance(value, Literal)
160

    
161
def is_null(value): return is_literal(value) and value.value == None
162

    
163
#### Composites
164

    
165
class List(Code):
166
    def __init__(self, values):
167
        Code.__init__(self)
168
        
169
        self.values = values
170
    
171
    def to_str(self, db): return ', '.join((v.to_str(db) for v in self.values))
172

    
173
class Tuple(List):
174
    def __init__(self, *values):
175
        List.__init__(self, values)
176
    
177
    def to_str(self, db): return '('+List.to_str(self, db)+')'
178

    
179
class Row(Tuple):
180
    def to_str(self, db): return 'ROW'+Tuple.to_str(self, db)
181

    
182
### Arrays
183

    
184
class Array(List):
185
    def __init__(self, values):
186
        values = map(remove_col_rename, values)
187
        
188
        List.__init__(self, values)
189
    
190
    def to_str(self, db): return 'ARRAY['+List.to_str(self, db)+']'
191

    
192
def to_Array(value):
193
    if isinstance(value, Array): return value
194
    return Array(lists.mk_seq(value))
195

    
196
##### Derived elements
197

    
198
src_self = object() # tells Col that it is its own source column
199

    
200
class Derived(Code):
201
    def __init__(self, srcs):
202
        '''An element which was derived from some other element(s).
203
        @param srcs See self.set_srcs()
204
        '''
205
        Code.__init__(self)
206
        
207
        self.set_srcs(srcs)
208
    
209
    def set_srcs(self, srcs, overwrite=True):
210
        '''
211
        @param srcs (self_type...)|src_self The element(s) this is derived from
212
        '''
213
        if not overwrite and self.srcs != (): return # already set
214
        
215
        if srcs == src_self: srcs = (self,)
216
        srcs = tuple(srcs) # make Col hashable
217
        self.srcs = srcs
218
    
219
    def _compare_on(self):
220
        compare_on = self.__dict__.copy()
221
        del compare_on['srcs'] # ignore
222
        return compare_on
223

    
224
def cols_srcs(cols): return lists.uniqify(iters.flatten((v.srcs for v in cols)))
225

    
226
##### Tables
227

    
228
class Table(Derived):
229
    def __init__(self, name, schema=None, srcs=(), is_temp=False):
230
        '''
231
        @param schema str|None (for no schema)
232
        @param srcs (Table...)|src_self See Derived.set_srcs()
233
        '''
234
        Derived.__init__(self, srcs)
235
        
236
        if util.is_str(name): name = truncate(name)
237
        
238
        self.name = name
239
        self.schema = schema
240
        self.is_temp = is_temp
241
        self.order_by = None
242
        self.index_cols = {}
243
    
244
    def to_str(self, db):
245
        str_ = ''
246
        if self.schema != None: str_ += as_Name(self.schema).to_str(db)+'.'
247
        str_ += as_Name(self.name).to_str(db)
248
        return str_
249
    
250
    def to_Table(self): return self
251
    
252
    def _compare_on(self):
253
        compare_on = Derived._compare_on(self)
254
        del compare_on['order_by'] # ignore
255
        del compare_on['index_cols'] # ignore
256
        return compare_on
257

    
258
def is_underlying_table(table):
259
    return isinstance(table, Table) and table.to_Table() is table
260

    
261
class NoUnderlyingTableException(Exception):
262
    def __init__(self, ref):
263
        Exception.__init__(self, 'for: '+strings.as_tt(strings.urepr(ref)))
264
        self.ref = ref
265

    
266
def underlying_table(table):
267
    table = remove_table_rename(table)
268
    if table != None and table.srcs:
269
        table, = table.srcs # for derived tables or row vars
270
    if not is_underlying_table(table): raise NoUnderlyingTableException(table)
271
    return table
272

    
273
def as_Table(table, schema=None):
274
    if table == None or isinstance(table, Code): return table
275
    else: return Table(table, schema)
276

    
277
def suffixed_table(table, suffix):
278
    table = copy.copy(table) # don't modify input!
279
    table.name = concat(table.name, suffix)
280
    return table
281

    
282
class NamedTable(Table):
283
    def __init__(self, name, code, cols=None):
284
        Table.__init__(self, name)
285
        
286
        code = as_Table(code)
287
        if not isinstance(code, (Table, FunctionCall, Expr)): code = Expr(code)
288
        if cols != None: cols = [to_name_only_col(c).to_Col() for c in cols]
289
        
290
        self.code = code
291
        self.cols = cols
292
    
293
    def to_str(self, db):
294
        str_ = self.code.to_str(db)
295
        if str_.find('\n') >= 0: whitespace = '\n'
296
        else: whitespace = ' '
297
        str_ += whitespace+'AS '+Table.to_str(self, db)
298
        if self.cols != None:
299
            str_ += ' ('+(', '.join((c.to_str(db) for c in self.cols)))+')'
300
        return str_
301
    
302
    def to_Table(self): return Table(self.name)
303

    
304
def remove_table_rename(table):
305
    if isinstance(table, NamedTable): table = table.code
306
    return table
307

    
308
##### Columns
309

    
310
class Col(Derived):
311
    def __init__(self, name, table=None, srcs=()):
312
        '''
313
        @param table Table|None (for no table)
314
        @param srcs (Col...)|src_self See Derived.set_srcs()
315
        '''
316
        Derived.__init__(self, srcs)
317
        
318
        if util.is_str(name): name = truncate(name)
319
        if util.is_str(table): table = Table(table)
320
        assert table == None or isinstance(table, Table)
321
        
322
        self.name = name
323
        self.table = table
324
    
325
    def to_str(self, db, for_str=False):
326
        str_ = as_Name(self.name).to_str(db)
327
        if for_str: str_ = clean_name(str_)
328
        if self.table != None:
329
            table = self.table.to_Table()
330
            if for_str: str_ = concat(strings.ustr(table), '.'+str_)
331
            else: str_ = table.to_str(db)+'.'+str_
332
        return str_
333
    
334
    def __str__(self): return self.to_str(mockDb, for_str=True)
335
    
336
    def to_Col(self): return self
337

    
338
def is_col(col): return isinstance(col, Col)
339

    
340
def is_table_col(col): return is_col(col) and col.table != None
341

    
342
def index_col(col):
343
    if not is_table_col(col): return None
344
    
345
    table = col.table
346
    try: name = table.index_cols[col.name]
347
    except KeyError: return None
348
    else: return Col(name, table, col.srcs)
349

    
350
def is_temp_col(col): return is_table_col(col) and col.table.is_temp
351

    
352
def as_Col(col, table=None, name=None):
353
    '''
354
    @param name If not None, any non-Col input will be renamed using NamedCol.
355
    '''
356
    if name != None:
357
        col = as_Value(col)
358
        if not isinstance(col, Col): col = NamedCol(name, col)
359
    
360
    if isinstance(col, Code): return col
361
    elif util.is_str(col): return Col(col, table)
362
    else: return Literal(col)
363

    
364
def with_table(col, table):
365
    if isinstance(col, NamedCol): pass # doesn't take a table
366
    elif isinstance(col, FunctionCall):
367
        col = copy.deepcopy(col) # don't modify input!
368
        col.args[0].table = table
369
    elif isinstance(col, Col):
370
        col = copy.copy(col) # don't modify input!
371
        col.table = table
372
    return col
373

    
374
def with_default_table(col, table):
375
    col = as_Col(col)
376
    if col.table == None: col = with_table(col, table)
377
    return col
378

    
379
def set_cols_table(table, cols):
380
    table = as_Table(table)
381
    
382
    for i, col in enumerate(cols):
383
        col = cols[i] = as_Col(col)
384
        col.table = table
385

    
386
def to_name_only_col(col, check_table=None):
387
    col = as_Col(col)
388
    if not is_table_col(col): return col
389
    
390
    if check_table != None:
391
        table = col.table
392
        assert table == None or table == check_table
393
    return Col(col.name)
394

    
395
def suffixed_col(col, suffix):
396
    return Col(concat(col.name, suffix), col.table, col.srcs)
397

    
398
def has_srcs(col): return is_col(col) and col.srcs
399

    
400
def cross_join_srcs(cols):
401
    cols = filter(has_srcs, cols) # empty srcs will mess up the cross join
402
    srcs = [[s.name for s in c.srcs] for c in cols]
403
    if not srcs: return [] # itertools.product() returns [()] for empty input
404
    return [Col(','.join(s)) for s in itertools.product(*srcs)]
405

    
406
class NamedCol(Col):
407
    def __init__(self, name, code):
408
        Col.__init__(self, name)
409
        
410
        code = as_Value(code)
411
        
412
        self.code = code
413
    
414
    def to_str(self, db):
415
        return self.code.to_str(db)+' AS '+Col.to_str(self, db)
416
    
417
    def to_Col(self): return Col(self.name)
418

    
419
def remove_col_rename(col):
420
    if isinstance(col, NamedCol): col = col.code
421
    return col
422

    
423
def underlying_col(col):
424
    col = remove_col_rename(col)
425
    if not isinstance(col, Col): raise NoUnderlyingTableException(col)
426
    
427
    return Col(col.name, underlying_table(col.table), col.srcs)
428

    
429
def wrap(wrap_func, value):
430
    '''Wraps a value, propagating any column renaming to the returned value.'''
431
    if isinstance(value, NamedCol):
432
        return NamedCol(value.name, wrap_func(value.code))
433
    else: return wrap_func(value)
434

    
435
class ColDict(dicts.DictProxy):
436
    '''A dict that automatically makes inserted entries Col objects.
437
    Anything that isn't a column is wrapped in a NamedCol with the key's column
438
    name by `as_Col(value, name=key.name)`.
439
    '''
440
    
441
    def __init__(self, db, keys_table, dict_={}):
442
        dicts.DictProxy.__init__(self, OrderedDict())
443
        
444
        keys_table = as_Table(keys_table)
445
        
446
        self.db = db
447
        self.table = keys_table
448
        self.update(dict_) # after setting vars because __setitem__() needs them
449
    
450
    def copy(self): return ColDict(self.db, self.table, self.inner.copy())
451
    
452
    def __getitem__(self, key):
453
        return dicts.DictProxy.__getitem__(self, self._key(key))
454
    
455
    def __setitem__(self, key, value):
456
        key = self._key(key)
457
        if value == None:
458
            try: value = self.db.col_info(key).default
459
            except NoUnderlyingTableException: pass # not a table column
460
        dicts.DictProxy.__setitem__(self, key, as_Col(value, name=key.name))
461
    
462
    def _key(self, key): return as_Col(key, self.table)
463

    
464
##### Definitions
465

    
466
class TypedCol(Col):
467
    def __init__(self, name, type_, default=None, nullable=True,
468
        constraints=None):
469
        assert default == None or isinstance(default, Code)
470
        
471
        Col.__init__(self, name)
472
        
473
        self.type = type_
474
        self.default = default
475
        self.nullable = nullable
476
        self.constraints = constraints
477
    
478
    def to_str(self, db):
479
        str_ = Col.to_str(self, db)+' '+as_Code(self.type).to_str(db)
480
        if not self.nullable: str_ += ' NOT NULL'
481
        if self.default != None: str_ += ' DEFAULT '+self.default.to_str(db)
482
        if self.constraints != None: str_ += ' '+self.constraints
483
        return str_
484
    
485
    def to_Col(self): return Col(self.name)
486

    
487
class SetOf(Code):
488
    def __init__(self, type_):
489
        Code.__init__(self)
490
        
491
        self.type = type_
492
    
493
    def to_str(self, db):
494
        return 'SETOF '+self.type.to_str(db)
495

    
496
class RowType(Code):
497
    def __init__(self, table):
498
        Code.__init__(self)
499
        
500
        self.table = table
501
    
502
    def to_str(self, db):
503
        return self.table.to_str(db)+'%ROWTYPE'
504

    
505
class ColType(Code):
506
    def __init__(self, col):
507
        Code.__init__(self)
508
        
509
        self.col = col
510
    
511
    def to_str(self, db):
512
        return self.col.to_str(db)+'%TYPE'
513

    
514
class ArrayType(Code):
515
    def __init__(self, elem_type):
516
        Code.__init__(self)
517
        elem_type = as_Code(elem_type)
518
        
519
        self.elem_type = elem_type
520
    
521
    def to_str(self, db):
522
        return self.elem_type.to_str(db)+'[]'
523

    
524
##### Functions
525

    
526
Function = Table
527
as_Function = as_Table
528

    
529
class InternalFunction(CustomCode): pass
530

    
531
#### Calls
532

    
533
class NamedArg(NamedCol):
534
    def __init__(self, name, value):
535
        NamedCol.__init__(self, name, value)
536
    
537
    def to_str(self, db):
538
        return Col.to_str(self, db)+' := '+self.code.to_str(db)
539

    
540
class FunctionCall(Code):
541
    def __init__(self, function, *args, **kw_args):
542
        '''
543
        @param args [Code|literal-value...] The function's arguments
544
        '''
545
        Code.__init__(self)
546
        
547
        function = as_Function(function)
548
        def filter_(arg): return remove_col_rename(as_Value(arg))
549
        args = map(filter_, args)
550
        args += [NamedArg(k, filter_(v)) for k, v in kw_args.iteritems()]
551
        
552
        self.function = function
553
        self.args = args
554
    
555
    def to_str(self, db):
556
        args_str = ', '.join((v.to_str(db) for v in self.args))
557
        return self.function.to_str(db)+'('+args_str+')'
558

    
559
def wrap_in_func(function, value):
560
    '''Wraps a value inside a function call.
561
    Propagates any column renaming to the returned value.
562
    '''
563
    return wrap(lambda v: FunctionCall(function, v), value)
564

    
565
def unwrap_func_call(func_call, check_name=None):
566
    '''Unwraps any function call to its first argument.
567
    Also removes any column renaming.
568
    '''
569
    func_call = remove_col_rename(func_call)
570
    if not isinstance(func_call, FunctionCall): return func_call
571
    
572
    if check_name != None:
573
        name = func_call.function.name
574
        assert name == None or name == check_name
575
    return func_call.args[0]
576

    
577
#### Definitions
578

    
579
class FunctionDef(Code):
580
    def __init__(self, function, return_type, body, params=[], modifiers=None):
581
        Code.__init__(self)
582
        
583
        return_type = as_Code(return_type)
584
        body = as_Code(body)
585
        
586
        self.function = function
587
        self.return_type = return_type
588
        self.body = body
589
        self.params = params
590
        self.modifiers = modifiers
591
    
592
    def to_str(self, db):
593
        params_str = (', '.join((p.to_str(db) for p in self.params)))
594
        str_ = '''\
595
CREATE FUNCTION '''+self.function.to_str(db)+'''('''+params_str+''')
596
RETURNS '''+self.return_type.to_str(db)+'''
597
LANGUAGE '''+self.body.lang+'''
598
'''
599
        if self.modifiers != None: str_ += self.modifiers+'\n'
600
        str_ += '''\
601
AS $$
602
'''+self.body.to_str(db)+'''
603
$$;
604
'''
605
        return str_
606

    
607
class FunctionParam(TypedCol):
608
    def __init__(self, name, type_, default=None, out=False):
609
        TypedCol.__init__(self, name, type_, default)
610
        
611
        self.out = out
612
    
613
    def to_str(self, db):
614
        str_ = TypedCol.to_str(self, db)
615
        if self.out: str_ = 'OUT '+str_
616
        return str_
617
    
618
    def to_Col(self): return Col(self.name)
619

    
620
### PL/pgSQL
621

    
622
class ReturnQuery(Code):
623
    def __init__(self, query):
624
        Code.__init__(self)
625
        
626
        query = as_Code(query)
627
        
628
        self.query = query
629
    
630
    def to_str(self, db):
631
        return 'RETURN QUERY\n'+strings.indent(self.query.to_str(db))+';\n'
632

    
633
## Exceptions
634

    
635
class BaseExcHandler(BasicObject):
636
    def to_str(self, db, body): raise NotImplementedError()
637
    
638
    def __repr__(self): return self.to_str(mockDb, '<body>')
639

    
640
suppress_exc = 'NULL;\n';
641

    
642
reraise_exc = 'RAISE USING ERRCODE = SQLSTATE, MESSAGE = SQLERRM;\n';
643

    
644
class ExcHandler(BaseExcHandler):
645
    def __init__(self, exc, handler=None):
646
        if handler != None: handler = as_Code(handler)
647
        
648
        self.exc = exc
649
        self.handler = handler
650
    
651
    def to_str(self, db, body):
652
        body = as_Code(body)
653
        
654
        if self.handler != None:
655
            handler_str = '\n'+strings.indent(self.handler.to_str(db), 2)
656
        else: handler_str = ' '+suppress_exc
657
        
658
        str_ = '''\
659
BEGIN
660
'''+strings.indent(body.to_str(db))+'''\
661
EXCEPTION
662
    WHEN '''+self.exc+''' THEN'''+handler_str+'''\
663
END;\
664
'''
665
        return str_
666

    
667
class NestedExcHandler(BaseExcHandler):
668
    def __init__(self, *handlers):
669
        '''
670
        @param handlers Sorted from outermost to innermost
671
        '''
672
        self.handlers = handlers
673
    
674
    def to_str(self, db, body):
675
        for handler in reversed(self.handlers): body = handler.to_str(db, body)
676
        return body
677

    
678
class ExcToWarning(Code):
679
    def __init__(self, return_):
680
        '''
681
        @param return_ Statement to return a default value in case of error
682
        '''
683
        Code.__init__(self)
684
        
685
        return_ = as_Code(return_)
686
        
687
        self.return_ = return_
688
    
689
    def to_str(self, db):
690
        return "RAISE WARNING '%', SQLERRM;\n"+self.return_.to_str(db)
691

    
692
unique_violation_handler = ExcHandler('unique_violation')
693

    
694
# Note doubled "\"s because inside Python string
695
plpythonu_error_handler = ExcHandler('internal_error', '''\
696
-- Handle PL/Python exceptions
697
DECLARE
698
    matches text[] := regexp_matches(SQLERRM,
699
        E'^(?:PL/Python: )?(\\\\w+): (.*)$'); -- .* also matches \\n
700
    exc_name text := matches[1];
701
    msg text := matches[2];
702
BEGIN
703
    /* Re-raise PL/Python exceptions with the PL/Python prefix removed.
704
    This allows the exception to be parsed like a native exception.
705
    Always raise as data_exception so it goes in the errors table. */
706
    IF exc_name IS NOT NULL THEN
707
        RAISE data_exception USING MESSAGE = msg;
708
    -- Re-raise non-PL/Python exceptions
709
    ELSE
710
        '''+reraise_exc+'''\
711
    END IF;
712
END;
713
''')
714

    
715
def data_exception_handler(handler):
716
    return ExcHandler('data_exception', handler)
717

    
718
row_var = Table('row')
719

    
720
class RowExcIgnore(Code):
721
    def __init__(self, row_type, select_query, with_row, cols=None,
722
        exc_handler=unique_violation_handler, row_var=row_var):
723
        '''
724
        @param row_type Ignored if a custom row_var is used.
725
        @pre If a custom row_var is used, it must already be defined.
726
        '''
727
        Code.__init__(self, lang='plpgsql')
728
        
729
        row_type = as_Code(row_type)
730
        select_query = as_Code(select_query)
731
        with_row = as_Code(with_row)
732
        row_var = as_Table(row_var)
733
        
734
        self.row_type = row_type
735
        self.select_query = select_query
736
        self.with_row = with_row
737
        self.cols = cols
738
        self.exc_handler = exc_handler
739
        self.row_var = row_var
740
    
741
    def to_str(self, db):
742
        if self.cols == None: row_vars = [self.row_var]
743
        else: row_vars = [Col(c.name, self.row_var) for c in self.cols]
744
        
745
        # Need an EXCEPTION block for each individual row because "When an error
746
        # is caught by an EXCEPTION clause, [...] all changes to persistent
747
        # database state within the block are rolled back."
748
        # This is unfortunate because "A block containing an EXCEPTION clause is
749
        # significantly more expensive to enter and exit than a block without
750
        # one."
751
        # (http://www.postgresql.org/docs/8.3/static/\
752
        # plpgsql-control-structures.html#PLPGSQL-ERROR-TRAPPING)
753
        str_ = '''\
754
FOR '''+(', '.join((v.to_str(db) for v in row_vars)))+''' IN
755
'''+strings.indent(self.select_query.to_str(db))+'''\
756
LOOP
757
'''+strings.indent(self.exc_handler.to_str(db, self.with_row))+'''\
758
END LOOP;
759
'''
760
        if self.row_var == row_var:
761
            str_ = '''\
762
DECLARE
763
    '''+self.row_var.to_str(db)+''' '''+self.row_type.to_str(db)+''';
764
BEGIN
765
'''+strings.indent(str_)+'''\
766
END;
767
'''
768
        return str_
769

    
770
##### Casts
771

    
772
class Cast(FunctionCall):
773
    def __init__(self, type_, value):
774
        type_ = as_Code(type_)
775
        value = as_Value(value)
776
        
777
        # Most types cannot be cast directly to unknown
778
        if type_.to_str(mockDb) == 'unknown': value = Cast('text', value)
779
        
780
        self.type_ = type_
781
        self.value = value
782
    
783
    def to_str(self, db):
784
        return 'CAST('+self.value.to_str(db)+' AS '+self.type_.to_str(db)+')'
785

    
786
def cast_literal(value):
787
    if not is_literal(value): return value
788
    
789
    if util.is_str(value.value): value = Cast('text', value)
790
    return value
791

    
792
##### Conditions
793

    
794
class NotCond(Code):
795
    def __init__(self, cond):
796
        Code.__init__(self)
797
        
798
        if not isinstance(cond, Coalesce): cond = Coalesce(cond, False)
799
        
800
        self.cond = cond
801
    
802
    def to_str(self, db): return 'NOT '+self.cond.to_str(db)
803

    
804
class ColValueCond(Code):
805
    def __init__(self, col, value):
806
        Code.__init__(self)
807
        
808
        value = as_ValueCond(value)
809
        
810
        self.col = col
811
        self.value = value
812
    
813
    def to_str(self, db): return self.value.to_str(db, self.col)
814

    
815
def combine_conds(conds, keyword=None):
816
    '''
817
    @param keyword The keyword to add before the conditions, if any
818
    '''
819
    str_ = ''
820
    if keyword != None:
821
        if conds == []: whitespace = ''
822
        elif len(conds) == 1: whitespace = ' '
823
        else: whitespace = '\n'
824
        str_ += keyword+whitespace
825
    
826
    str_ += '\nAND '.join(conds)
827
    return str_
828

    
829
##### Condition column comparisons
830

    
831
class ValueCond(BasicObject):
832
    def __init__(self, value):
833
        value = remove_col_rename(as_Value(value))
834
        
835
        self.value = value
836
    
837
    def to_str(self, db, left_value):
838
        '''
839
        @param left_value The Code object that the condition is being applied on
840
        '''
841
        raise NotImplemented()
842
    
843
    def __repr__(self): return self.to_str(mockDb, '<left_value>')
844

    
845
class CompareCond(ValueCond):
846
    def __init__(self, value, operator='='):
847
        '''
848
        @param operator By default, compares NULL values literally. Use '~=' or
849
            '~!=' to pass NULLs through.
850
        '''
851
        ValueCond.__init__(self, value)
852
        self.operator = operator
853
    
854
    def to_str(self, db, left_value):
855
        left_value = remove_col_rename(as_Col(left_value))
856
        
857
        right_value = self.value
858
        
859
        # Parse operator
860
        operator = self.operator
861
        passthru_null_ref = [False]
862
        operator = strings.remove_prefix('~', operator, passthru_null_ref)
863
        neg_ref = [False]
864
        operator = strings.remove_prefix('!', operator, neg_ref)
865
        equals = operator.endswith('=') # also includes <=, >=
866
        
867
        # Handle nullable columns
868
        check_null = False
869
        if not passthru_null_ref[0]: # NULLs compare equal
870
            try: left_value = ensure_not_null(db, left_value)
871
            except ensure_not_null_excs: # fall back to alternate method
872
                check_null = equals and isinstance(right_value, Col)
873
            else:
874
                if isinstance(left_value, EnsureNotNull):
875
                    right_value = ensure_not_null(db, right_value,
876
                        left_value.type) # apply same function to both sides
877
        
878
        if equals and is_null(right_value): operator = 'IS'
879
        
880
        left = left_value.to_str(db)
881
        right = right_value.to_str(db)
882
        
883
        # Create str
884
        str_ = left+' '+operator+' '+right
885
        if check_null:
886
            str_ = '('+str_+' OR ('+left+' IS NULL AND '+right+' IS NULL))'
887
        if neg_ref[0]: str_ = 'NOT '+str_
888
        return str_
889

    
890
# Tells as_ValueCond() to assume a non-ValueCond is a literal value
891
assume_literal = object()
892

    
893
def as_ValueCond(value, default_table=assume_literal):
894
    if not isinstance(value, ValueCond):
895
        if default_table is not assume_literal:
896
            value = with_default_table(value, default_table)
897
        return CompareCond(value)
898
    else: return value
899

    
900
##### Joins
901

    
902
join_same = object() # tells Join the left and right columns have the same name
903

    
904
# Tells Join the left and right columns have the same name and are never NULL
905
join_same_not_null = object()
906

    
907
filter_out = object() # tells Join to filter out rows that match the join
908

    
909
class Join(BasicObject):
910
    def __init__(self, table, mapping={}, type_=None):
911
        '''
912
        @param mapping dict(right_table_col=left_table_col, ...)
913
            or [using_col...]
914
            * if left_table_col is join_same: left_table_col = right_table_col
915
              * Note that right_table_col must be a string
916
            * if left_table_col is join_same_not_null:
917
              left_table_col = right_table_col and both have NOT NULL constraint
918
              * Note that right_table_col must be a string
919
        @param type_ None (for plain join)|str (e.g. 'LEFT')|filter_out
920
            * filter_out: equivalent to 'LEFT' with the query filtered by
921
              `table_pkey IS NULL` (indicating no match)
922
        '''
923
        if util.is_str(table): table = Table(table)
924
        if lists.is_seq(mapping):
925
            mapping = dict(((c, join_same_not_null) for c in mapping))
926
        assert type_ == None or util.is_str(type_) or type_ is filter_out
927
        
928
        self.table = table
929
        self.mapping = mapping
930
        self.type_ = type_
931
    
932
    def to_str(self, db, left_table_):
933
        def join(entry):
934
            '''Parses non-USING joins'''
935
            right_table_col, left_table_col = entry
936
            
937
            # Switch order (right_table_col is on the left in the comparison)
938
            left = right_table_col
939
            right = left_table_col
940
            left_table = self.table
941
            right_table = left_table_
942
            
943
            # Parse left side
944
            left = with_default_table(left, left_table)
945
            
946
            # Parse special values
947
            left_on_right = Col(left.name, right_table)
948
            if right is join_same: right = left_on_right
949
            elif right is join_same_not_null:
950
                right = CompareCond(left_on_right, '~=')
951
            
952
            # Parse right side
953
            right = as_ValueCond(right, right_table)
954
            
955
            return right.to_str(db, left)
956
        
957
        # Create join condition
958
        type_ = self.type_
959
        joins = self.mapping
960
        if joins == {}: join_cond = None
961
        elif type_ is not filter_out and reduce(operator.and_,
962
            (v is join_same_not_null for v in joins.itervalues())):
963
            # all cols w/ USING, so can use simpler USING syntax
964
            cols = map(to_name_only_col, joins.iterkeys())
965
            join_cond = 'USING ('+(', '.join((c.to_str(db) for c in cols)))+')'
966
        else: join_cond = combine_conds(map(join, joins.iteritems()), 'ON')
967
        
968
        if isinstance(self.table, NamedTable): whitespace = '\n'
969
        else: whitespace = ' '
970
        
971
        # Create join
972
        if type_ is filter_out: type_ = 'LEFT'
973
        str_ = ''
974
        if type_ != None: str_ += type_+' '
975
        str_ += 'JOIN'+whitespace+self.table.to_str(db)
976
        if join_cond != None: str_ += whitespace+join_cond
977
        return str_
978
    
979
    def __repr__(self): return self.to_str(mockDb, '<left_table>')
980

    
981
##### Value exprs
982

    
983
all_cols = CustomCode('*')
984

    
985
default = CustomCode('DEFAULT')
986

    
987
row_count = FunctionCall(InternalFunction('COUNT'), all_cols)
988

    
989
class Coalesce(FunctionCall):
990
    def __init__(self, *args):
991
        FunctionCall.__init__(self, InternalFunction('COALESCE'), *args)
992

    
993
class Nullif(FunctionCall):
994
    def __init__(self, *args):
995
        FunctionCall.__init__(self, InternalFunction('NULLIF'), *args)
996

    
997
null = Literal(None)
998
null_as_str = Cast('text', null)
999

    
1000
def to_text(value): return Coalesce(Cast('text', value), null_as_str)
1001

    
1002
# See <http://www.postgresql.org/docs/8.3/static/datatype-numeric.html>
1003
null_sentinels = {
1004
    'character varying': r'\N',
1005
    'double precision': 'NaN',
1006
    'integer': 2147483647,
1007
    'text': r'\N',
1008
    'date': 'infinity',
1009
    'timestamp with time zone': 'infinity',
1010
    'taxonrank': 'unknown',
1011
}
1012

    
1013
class EnsureNotNull(Coalesce):
1014
    def __init__(self, value, type_):
1015
        if isinstance(type_, ArrayType): null = []
1016
        else: null = null_sentinels[type_]
1017
        Coalesce.__init__(self, as_Col(value), Cast(type_, null))
1018
        
1019
        self.type = type_
1020
    
1021
    def to_str(self, db):
1022
        col = self.args[0]
1023
        index_col_ = index_col(col)
1024
        if index_col_ != None: return index_col_.to_str(db)
1025
        return Coalesce.to_str(self, db)
1026

    
1027
#### Arrays
1028

    
1029
class ArrayMerge(FunctionCall):
1030
    def __init__(self, sep, array):
1031
        array = to_Array(array)
1032
        FunctionCall.__init__(self, InternalFunction('array_to_string'), array,
1033
            sep)
1034

    
1035
def merge_not_null(db, sep, values):
1036
    return ArrayMerge(sep, map(to_text, values))
1037

    
1038
##### Table exprs
1039

    
1040
class Values(Code):
1041
    def __init__(self, values):
1042
        '''
1043
        @param values [...]|[[...], ...] Can be one or multiple rows.
1044
        '''
1045
        Code.__init__(self)
1046
        
1047
        rows = values
1048
        if len(values) >= 1 and not lists.is_seq(values[0]): # only one row
1049
            rows = [values]
1050
        for i, row in enumerate(rows):
1051
            rows[i] = map(remove_col_rename, map(as_Value, row))
1052
        
1053
        self.rows = rows
1054
    
1055
    def to_str(self, db):
1056
        return 'VALUES '+(', '.join((Tuple(*r).to_str(db) for r in self.rows)))
1057

    
1058
def NamedValues(name, cols, values):
1059
    '''
1060
    @param cols None|[...]
1061
    @post `cols` will be changed to Col objects with the table set to `name`.
1062
    '''
1063
    table = NamedTable(name, Values(values), cols)
1064
    if cols != None: set_cols_table(table, cols)
1065
    return table
1066

    
1067
##### Database structure
1068

    
1069
def is_nullable(db, value):
1070
    if not is_table_col(value): return is_null(value)
1071
    try: return db.col_info(value).nullable
1072
    except NoUnderlyingTableException: return True # not a table column
1073

    
1074
text_types = set(['character varying', 'text'])
1075

    
1076
def is_text_type(type_): return type_ in text_types
1077

    
1078
def is_text_col(db, col): return is_text_type(db.col_info(col).type)
1079

    
1080
def canon_type(type_):
1081
    if type_ in text_types: return 'text'
1082
    else: return type_
1083

    
1084
ensure_not_null_excs = (NoUnderlyingTableException, KeyError)
1085

    
1086
def ensure_not_null(db, col, type_=None):
1087
    '''
1088
    @param col If type_ is not set, must have an underlying column.
1089
    @param type_ If set, overrides the underlying column's type and casts the
1090
        column to it if needed.
1091
    @return EnsureNotNull|Col
1092
    @throws ensure_not_null_excs
1093
    '''
1094
    col = remove_col_rename(col)
1095
    
1096
    try: col_type = db.col_info(underlying_col(col)).type
1097
    except NoUnderlyingTableException:
1098
        if type_ == None and is_null(col): raise # NULL has no type
1099
    else:
1100
        if type_ == None: type_ = col_type
1101
        elif type_ != col_type: col = Cast(type_, col)
1102
    
1103
    if is_nullable(db, col):
1104
        try: col = EnsureNotNull(col, type_)
1105
        except KeyError, e:
1106
            # Warn of no null sentinel for type, even if caller catches error
1107
            warnings.warn(UserWarning(exc.str_(e)))
1108
            raise
1109
    
1110
    return col
1111

    
1112
def try_mk_not_null(db, value):
1113
    '''
1114
    Warning: This function does not guarantee that its result is NOT NULL.
1115
    '''
1116
    try: return ensure_not_null(db, value)
1117
    except ensure_not_null_excs: return value
1118

    
1119
##### Expression transforming
1120

    
1121
true_expr = 'true'
1122
false_expr = 'false'
1123

    
1124
true_re = true_expr
1125
false_re = false_expr
1126
bool_re = r'(?:'+true_re+r'|'+false_re+r')'
1127
atom_re = r'(?:'+bool_re+r'|\([^()]*\)'+r')'
1128

    
1129
def logic_op_re(op, value_re, expr_re=''):
1130
    op_re = ' '+op+' '
1131
    return '(?:'+expr_re+op_re+value_re+'|'+value_re+op_re+expr_re+')'
1132

    
1133
not_re = r'\bNOT '
1134
not_false_re = not_re+false_re+r'\b'
1135
not_true_re = not_re+true_re+r'\b'
1136
and_false_re = logic_op_re('AND', false_re, atom_re)
1137
and_false_not_true_re = '(?:'+not_true_re+'|'+and_false_re+')'
1138
and_true_re = logic_op_re('AND', true_re)
1139
or_re = logic_op_re('OR', bool_re)
1140
or_and_true_re = '(?:'+and_true_re+'|'+or_re+')'
1141

    
1142
def simplify_parens(expr):
1143
    return regexp.sub_nested(r'\(('+atom_re+')\)', r'\1', expr)
1144

    
1145
def simplify_recursive(sub_func, expr):
1146
    '''
1147
    @param sub_func See regexp.sub_recursive() sub_func param
1148
    '''
1149
    return regexp.sub_recursive(lambda s: sub_func(simplify_parens(s)), expr)
1150
        # simplify_parens() is also done at end in final iteration
1151

    
1152
def simplify_expr(expr):
1153
    def simplify_logic_ops(expr):
1154
        total_n = 0
1155
        expr, n = re.subn(not_false_re, true_re, expr)
1156
        total_n += n
1157
        expr, n = re.subn(and_false_not_true_re, false_expr, expr)
1158
        total_n += n
1159
        expr, n = re.subn(or_and_true_re, r'', expr)
1160
        total_n += n
1161
        return expr, total_n
1162
    
1163
    expr = expr.replace('NULL IS NULL', true_expr)
1164
    expr = expr.replace('NULL IS NOT NULL', false_expr)
1165
    expr = simplify_recursive(simplify_logic_ops, expr)
1166
    return expr
1167

    
1168
name_re = r'(?:\w+|(?:"[^"]*")+)'
1169

    
1170
def parse_expr_col(str_):
1171
    match = re.match(r'^\('+name_re+r'\(('+name_re+r').*\)\)$', str_)
1172
    if match: str_ = match.group(1)
1173
    return unesc_name(str_)
1174

    
1175
def map_expr(db, expr, mapping, in_cols_found=None):
1176
    '''Replaces output columns with input columns in an expression.
1177
    @param in_cols_found If set, will be filled in with the expr's (input) cols
1178
    '''
1179
    for out, in_ in mapping.iteritems():
1180
        orig_expr = expr
1181
        out = to_name_only_col(out)
1182
        in_str = to_name_only_col(remove_col_rename(in_)).to_str(db)
1183
        
1184
        # Replace out both with and without quotes
1185
        expr = expr.replace(out.to_str(db), in_str)
1186
        expr = re.sub(r'(?<!["\'\.=\[])\b'+out.name+r'\b(?!["\',\.=\]])',
1187
            in_str, expr)
1188
        
1189
        if in_cols_found != None and expr != orig_expr: # replaced something
1190
            in_cols_found.append(in_)
1191
    
1192
    return simplify_expr(expr)
(37-37/51)