Project

General

Profile

1
# SQL code generation
2

    
3
import copy
4
import operator
5
import re
6
import UserDict
7
import warnings
8

    
9
import dicts
10
import exc
11
import iters
12
import lists
13
import objects
14
import strings
15
import util
16

    
17
##### Names
18

    
19
identifier_max_len = 63 # works for both PostgreSQL and MySQL
20

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

    
32
def truncate(str_): return concat(str_, '')
33

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

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

    
46
def clean_name(name): return name.replace('"', '').replace('`', '')
47

    
48
def esc_comment(comment): return '/*'+comment.replace('*/', '* /')+'*/'
49

    
50
##### General SQL code objects
51

    
52
class MockDb:
53
    def esc_value(self, value): return strings.repr_no_u(value)
54
    
55
    def esc_name(self, name): return esc_name(name)
56
    
57
    def col_info(self, col):
58
        return TypedCol(col.name, '<type>', CustomCode('<default>'), True)
59

    
60
mockDb = MockDb()
61

    
62
class BasicObject(objects.BasicObject):
63
    def __init__(self, value): self.value = value
64
    
65
    def __str__(self): return clean_name(strings.repr_no_u(self))
66

    
67
##### Unparameterized code objects
68

    
69
class Code(BasicObject):
70
    def to_str(self, db): raise NotImplementedError()
71
    
72
    def __repr__(self): return self.to_str(mockDb)
73

    
74
class CustomCode(Code):
75
    def __init__(self, str_): self.str_ = str_
76
    
77
    def to_str(self, db): return self.str_
78

    
79
def as_Code(value, db=None):
80
    '''
81
    @param db If set, runs db.std_code() on the value.
82
    '''
83
    if util.is_str(value):
84
        if db != None: value = db.std_code(value)
85
        return CustomCode(value)
86
    else: return Literal(value)
87

    
88
class Expr(Code):
89
    def __init__(self, expr): self.expr = expr
90
    
91
    def to_str(self, db): return '('+self.expr.to_str(db)+')'
92

    
93
##### Literal values
94

    
95
class Literal(Code):
96
    def __init__(self, value): self.value = value
97
    
98
    def to_str(self, db): return db.esc_value(self.value)
99

    
100
def as_Value(value):
101
    if isinstance(value, Code): return value
102
    else: return Literal(value)
103

    
104
def is_null(value): return isinstance(value, Literal) and value.value == None
105

    
106
##### Derived elements
107

    
108
src_self = object() # tells Col that it is its own source column
109

    
110
class Derived(Code):
111
    def __init__(self, srcs):
112
        '''An element which was derived from some other element(s).
113
        @param srcs See self.set_srcs()
114
        '''
115
        self.set_srcs(srcs)
116
    
117
    def set_srcs(self, srcs, overwrite=True):
118
        '''
119
        @param srcs (self_type...)|src_self The element(s) this is derived from
120
        '''
121
        if not overwrite and self.srcs != (): return # already set
122
        
123
        if srcs == src_self: srcs = (self,)
124
        srcs = tuple(srcs) # make Col hashable
125
        self.srcs = srcs
126
    
127
    def _compare_on(self):
128
        compare_on = self.__dict__.copy()
129
        del compare_on['srcs'] # ignore
130
        return compare_on
131

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

    
134
##### Tables
135

    
136
class Table(Derived):
137
    def __init__(self, name, schema=None, srcs=(), is_temp=False):
138
        '''
139
        @param schema str|None (for no schema)
140
        @param srcs (Table...)|src_self See Derived.set_srcs()
141
        '''
142
        Derived.__init__(self, srcs)
143
        
144
        name = truncate(name)
145
        
146
        self.name = name
147
        self.schema = schema
148
        self.is_temp = is_temp
149
        self.index_cols = {}
150
    
151
    def to_str(self, db):
152
        str_ = ''
153
        if self.schema != None: str_ += db.esc_name(self.schema)+'.'
154
        str_ += db.esc_name(self.name)
155
        return str_
156
    
157
    def to_Table(self): return self
158
    
159
    def _compare_on(self):
160
        compare_on = Derived._compare_on(self)
161
        del compare_on['index_cols'] # ignore
162
        return compare_on
163

    
164
def is_underlying_table(table):
165
    return isinstance(table, Table) and table.to_Table() is table
166

    
167
class NoUnderlyingTableException(Exception): pass
168

    
169
def underlying_table(table):
170
    table = remove_table_rename(table)
171
    if not is_underlying_table(table): raise NoUnderlyingTableException
172
    return table
173

    
174
def as_Table(table, schema=None):
175
    if table == None or isinstance(table, Code): return table
176
    else: return Table(table, schema)
177

    
178
def suffixed_table(table, suffix): return Table(table.name+suffix, table.schema)
179

    
180
class NamedTable(Table):
181
    def __init__(self, name, code, cols=None):
182
        Table.__init__(self, name)
183
        
184
        code = as_Table(code)
185
        if not isinstance(code, (Table, FunctionCall, Expr)): code = Expr(code)
186
        if cols != None: cols = [to_name_only_col(c).to_Col() for c in cols]
187
        
188
        self.code = code
189
        self.cols = cols
190
    
191
    def to_str(self, db):
192
        str_ = self.code.to_str(db)
193
        if str_.find('\n') >= 0: whitespace = '\n'
194
        else: whitespace = ' '
195
        str_ += whitespace+'AS '+Table.to_str(self, db)
196
        if self.cols != None:
197
            str_ += ' ('+(', '.join((c.to_str(db) for c in self.cols)))+')'
198
        return str_
199
    
200
    def to_Table(self): return Table(self.name)
201

    
202
def remove_table_rename(table):
203
    if isinstance(table, NamedTable): table = table.code
204
    return table
205

    
206
##### Columns
207

    
208
class Col(Derived):
209
    def __init__(self, name, table=None, srcs=()):
210
        '''
211
        @param table Table|None (for no table)
212
        @param srcs (Col...)|src_self See Derived.set_srcs()
213
        '''
214
        Derived.__init__(self, srcs)
215
        
216
        name = truncate(name)
217
        if util.is_str(table): table = Table(table)
218
        assert table == None or isinstance(table, Table)
219
        
220
        self.name = name
221
        self.table = table
222
    
223
    def to_str(self, db, for_str=False):
224
        str_ = db.esc_name(self.name)
225
        if for_str: str_ = clean_name(str_)
226
        if self.table != None:
227
            table = self.table.to_Table()
228
            if for_str: str_ = concat(str(table), '.'+str_)
229
            else: str_ = table.to_str(db)+'.'+str_
230
        return str_
231
    
232
    def __str__(self): return self.to_str(mockDb, for_str=True)
233
    
234
    def to_Col(self): return self
235

    
236
def is_table_col(col): return isinstance(col, Col) and col.table != None
237

    
238
def index_col(col):
239
    if not is_table_col(col): return None
240
    return col.table.index_cols.get(col.name, None)
241

    
242
def is_temp_col(col): return is_table_col(col) and col.table.is_temp
243

    
244
def as_Col(col, table=None, name=None):
245
    '''
246
    @param name If not None, any non-Col input will be renamed using NamedCol.
247
    '''
248
    if name != None:
249
        col = as_Value(col)
250
        if not isinstance(col, Col): col = NamedCol(name, col)
251
    
252
    if isinstance(col, Code): return col
253
    else: return Col(col, table)
254

    
255
def with_default_table(col, table, overwrite=False):
256
    col = as_Col(col)
257
    if not isinstance(col, NamedCol) and (overwrite or col.table == None):
258
        col = copy.copy(col) # don't modify input!
259
        col.table = table
260
    return col
261

    
262
def set_cols_table(table, cols):
263
    table = as_Table(table)
264
    
265
    for i, col in enumerate(cols):
266
        col = cols[i] = as_Col(col)
267
        col.table = table
268

    
269
def to_name_only_col(col, check_table=None):
270
    col = as_Col(col)
271
    if not is_table_col(col): return col
272
    
273
    if check_table != None:
274
        table = col.table
275
        assert table == None or table == check_table
276
    return Col(col.name)
277

    
278
def suffixed_col(col, suffix):
279
    return Col(concat(col.name, suffix), col.table, col.srcs)
280

    
281
class NamedCol(Col):
282
    def __init__(self, name, code):
283
        Col.__init__(self, name)
284
        
285
        code = as_Value(code)
286
        
287
        self.code = code
288
    
289
    def to_str(self, db):
290
        return self.code.to_str(db)+' AS '+Col.to_str(self, db)
291
    
292
    def to_Col(self): return Col(self.name)
293

    
294
def remove_col_rename(col):
295
    if isinstance(col, NamedCol): col = col.code
296
    return col
297

    
298
def underlying_col(col):
299
    col = remove_col_rename(col)
300
    if not isinstance(col, Col): raise NoUnderlyingTableException
301
    
302
    return Col(col.name, underlying_table(col.table), col.srcs)
303

    
304
def wrap(wrap_func, value):
305
    '''Wraps a value, propagating any column renaming to the returned value.'''
306
    if isinstance(value, NamedCol):
307
        return NamedCol(value.name, wrap_func(value.code))
308
    else: return wrap_func(value)
309

    
310
class ColDict(dicts.DictProxy):
311
    '''A dict that automatically makes inserted entries Col objects'''
312
    
313
    def __init__(self, db, keys_table, dict_={}):
314
        dicts.DictProxy.__init__(self, {})
315
        
316
        keys_table = as_Table(keys_table)
317
        
318
        self.db = db
319
        self.table = keys_table
320
        self.update(dict_) # after setting vars because __setitem__() needs them
321
    
322
    def copy(self): return ColDict(self.db, self.table, self.inner.copy())
323
    
324
    def __getitem__(self, key):
325
        return dicts.DictProxy.__getitem__(self, self._key(key))
326
    
327
    def __setitem__(self, key, value):
328
        key = self._key(key)
329
        if value == None: value = self.db.col_info(key).default
330
        dicts.DictProxy.__setitem__(self, key, as_Col(value, name=key.name))
331
    
332
    def _key(self, key): return as_Col(key, self.table)
333

    
334
##### Functions
335

    
336
Function = Table
337
as_Function = as_Table
338

    
339
class InternalFunction(CustomCode): pass
340

    
341
class NamedArg(NamedCol):
342
    def __init__(self, name, value):
343
        NamedCol.__init__(self, name, value)
344
    
345
    def to_str(self, db):
346
        return Col.to_str(self, db)+' := '+self.code.to_str(db)
347

    
348
class FunctionCall(Code):
349
    def __init__(self, function, *args, **kw_args):
350
        '''
351
        @param args [Code|literal-value...] The function's arguments
352
        '''
353
        function = as_Function(function)
354
        def filter_(arg): return remove_col_rename(as_Value(arg))
355
        args = map(filter_, args)
356
        args += [NamedArg(k, filter_(v)) for k, v in kw_args.iteritems()]
357
        
358
        self.function = function
359
        self.args = args
360
    
361
    def to_str(self, db):
362
        args_str = ', '.join((v.to_str(db) for v in self.args))
363
        return self.function.to_str(db)+'('+args_str+')'
364

    
365
def wrap_in_func(function, value):
366
    '''Wraps a value inside a function call.
367
    Propagates any column renaming to the returned value.
368
    '''
369
    return wrap(lambda v: FunctionCall(function, v), value)
370

    
371
def unwrap_func_call(func_call, check_name=None):
372
    '''Unwraps any function call to its first argument.
373
    Also removes any column renaming.
374
    '''
375
    func_call = remove_col_rename(func_call)
376
    if not isinstance(func_call, FunctionCall): return func_call
377
    
378
    if check_name != None:
379
        name = func_call.function.name
380
        assert name == None or name == check_name
381
    return func_call.args[0]
382

    
383
##### Casts
384

    
385
class Cast(FunctionCall):
386
    def __init__(self, type_, value):
387
        value = as_Value(value)
388
        
389
        self.type_ = type_
390
        self.value = value
391
    
392
    def to_str(self, db):
393
        return 'CAST('+self.value.to_str(db)+' AS '+self.type_+')'
394

    
395
##### Conditions
396

    
397
class ColValueCond(Code):
398
    def __init__(self, col, value):
399
        value = as_ValueCond(value)
400
        
401
        self.col = col
402
        self.value = value
403
    
404
    def to_str(self, db): return self.value.to_str(db, self.col)
405

    
406
def combine_conds(conds, keyword=None):
407
    '''
408
    @param keyword The keyword to add before the conditions, if any
409
    '''
410
    str_ = ''
411
    if keyword != None:
412
        if conds == []: whitespace = ''
413
        elif len(conds) == 1: whitespace = ' '
414
        else: whitespace = '\n'
415
        str_ += keyword+whitespace
416
    
417
    str_ += '\nAND '.join(conds)
418
    return str_
419

    
420
##### Condition column comparisons
421

    
422
class ValueCond(BasicObject):
423
    def __init__(self, value):
424
        value = remove_col_rename(as_Value(value))
425
        
426
        self.value = value
427
    
428
    def to_str(self, db, left_value):
429
        '''
430
        @param left_value The Code object that the condition is being applied on
431
        '''
432
        raise NotImplemented()
433
    
434
    def __repr__(self): return self.to_str(mockDb, '<left_value>')
435

    
436
class CompareCond(ValueCond):
437
    def __init__(self, value, operator='='):
438
        '''
439
        @param operator By default, compares NULL values literally. Use '~=' or
440
            '~!=' to pass NULLs through.
441
        '''
442
        ValueCond.__init__(self, value)
443
        self.operator = operator
444
    
445
    def to_str(self, db, left_value):
446
        left_value = remove_col_rename(as_Col(left_value))
447
        
448
        right_value = self.value
449
        
450
        # Parse operator
451
        operator = self.operator
452
        passthru_null_ref = [False]
453
        operator = strings.remove_prefix('~', operator, passthru_null_ref)
454
        neg_ref = [False]
455
        operator = strings.remove_prefix('!', operator, neg_ref)
456
        equals = operator.endswith('=') # also includes <=, >=
457
        
458
        # Handle nullable columns
459
        check_null = False
460
        if not passthru_null_ref[0]: # NULLs compare equal
461
            try: left_value = ensure_not_null(db, left_value)
462
            except ensure_not_null_excs: # fall back to alternate method
463
                check_null = equals and isinstance(right_value, Col)
464
            else:
465
                if isinstance(left_value, EnsureNotNull):
466
                    right_value = ensure_not_null(db, right_value,
467
                        left_value.type) # apply same function to both sides
468
        
469
        if equals and is_null(right_value): operator = 'IS'
470
        
471
        left = left_value.to_str(db)
472
        right = right_value.to_str(db)
473
        
474
        # Create str
475
        str_ = left+' '+operator+' '+right
476
        if check_null:
477
            str_ = '('+str_+' OR ('+left+' IS NULL AND '+right+' IS NULL))'
478
        if neg_ref[0]: str_ = 'NOT '+str_
479
        return str_
480

    
481
# Tells as_ValueCond() to assume a non-ValueCond is a literal value
482
assume_literal = object()
483

    
484
def as_ValueCond(value, default_table=assume_literal):
485
    if not isinstance(value, ValueCond):
486
        if default_table is not assume_literal:
487
            value = with_default_table(value, default_table)
488
        return CompareCond(value)
489
    else: return value
490

    
491
##### Joins
492

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

    
495
# Tells Join the left and right columns have the same name and are never NULL
496
join_same_not_null = object()
497

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

    
500
class Join(BasicObject):
501
    def __init__(self, table, mapping={}, type_=None):
502
        '''
503
        @param mapping dict(right_table_col=left_table_col, ...)
504
            * if left_table_col is join_same: left_table_col = right_table_col
505
              * Note that right_table_col must be a string
506
            * if left_table_col is join_same_not_null:
507
              left_table_col = right_table_col and both have NOT NULL constraint
508
              * Note that right_table_col must be a string
509
        @param type_ None (for plain join)|str (e.g. 'LEFT')|filter_out
510
            * filter_out: equivalent to 'LEFT' with the query filtered by
511
              `table_pkey IS NULL` (indicating no match)
512
        '''
513
        if util.is_str(table): table = Table(table)
514
        assert type_ == None or util.is_str(type_) or type_ is filter_out
515
        
516
        self.table = table
517
        self.mapping = mapping
518
        self.type_ = type_
519
    
520
    def to_str(self, db, left_table_):
521
        def join(entry):
522
            '''Parses non-USING joins'''
523
            right_table_col, left_table_col = entry
524
            
525
            # Switch order (right_table_col is on the left in the comparison)
526
            left = right_table_col
527
            right = left_table_col
528
            left_table = self.table
529
            right_table = left_table_
530
            
531
            # Parse left side
532
            left = with_default_table(left, left_table)
533
            
534
            # Parse special values
535
            left_on_right = Col(left.name, right_table)
536
            if right is join_same: right = left_on_right
537
            elif right is join_same_not_null:
538
                right = CompareCond(left_on_right, '~=')
539
            
540
            # Parse right side
541
            right = as_ValueCond(right, right_table)
542
            
543
            return right.to_str(db, left)
544
        
545
        # Create join condition
546
        type_ = self.type_
547
        joins = self.mapping
548
        if joins == {}: join_cond = None
549
        elif type_ is not filter_out and reduce(operator.and_,
550
            (v is join_same_not_null for v in joins.itervalues())):
551
            # all cols w/ USING, so can use simpler USING syntax
552
            cols = map(to_name_only_col, joins.iterkeys())
553
            join_cond = 'USING ('+(', '.join((c.to_str(db) for c in cols)))+')'
554
        else: join_cond = combine_conds(map(join, joins.iteritems()), 'ON')
555
        
556
        if isinstance(self.table, NamedTable): whitespace = '\n'
557
        else: whitespace = ' '
558
        
559
        # Create join
560
        if type_ is filter_out: type_ = 'LEFT'
561
        str_ = ''
562
        if type_ != None: str_ += type_+' '
563
        str_ += 'JOIN'+whitespace+self.table.to_str(db)
564
        if join_cond != None: str_ += whitespace+join_cond
565
        return str_
566
    
567
    def __repr__(self): return self.to_str(mockDb, '<left_table>')
568

    
569
##### Value exprs
570

    
571
default = CustomCode('DEFAULT')
572

    
573
row_count = CustomCode('count(*)')
574

    
575
class Coalesce(FunctionCall):
576
    def __init__(self, *args):
577
        FunctionCall.__init__(self, InternalFunction('COALESCE'), *args)
578

    
579
class Nullif(FunctionCall):
580
    def __init__(self, *args):
581
        FunctionCall.__init__(self, InternalFunction('NULLIF'), *args)
582

    
583
# See <http://www.postgresql.org/docs/8.3/static/datatype-numeric.html>
584
null_sentinels = {
585
    'character varying': r'\N',
586
    'double precision': 'NaN',
587
    'integer': 2147483647,
588
    'text': r'\N',
589
    'timestamp with time zone': 'infinity'
590
}
591

    
592
class EnsureNotNull(Coalesce):
593
    def __init__(self, value, type_):
594
        Coalesce.__init__(self, as_Col(value),
595
            Cast(type_, null_sentinels[type_]))
596
        
597
        self.type = type_
598
    
599
    def to_str(self, db):
600
        col = self.args[0]
601
        index_col_ = index_col(col)
602
        if index_col_ != None: return index_col_.to_str(db)
603
        return Coalesce.to_str(self, db)
604

    
605
##### Table exprs
606

    
607
class Values(Code):
608
    def __init__(self, values):
609
        '''
610
        @param values [...]|[[...], ...] Can be one or multiple rows.
611
        '''
612
        rows = values
613
        if len(values) >= 1 and not lists.is_seq(values[0]): # only one row
614
            rows = [values]
615
        for i, row in enumerate(rows):
616
            rows[i] = map(remove_col_rename, map(as_Value, row))
617
        
618
        self.rows = rows
619
    
620
    def to_str(self, db):
621
        def row_str(row):
622
            return '('+(', '.join((v.to_str(db) for v in row)))+')'
623
        return 'VALUES '+(', '.join(map(row_str, self.rows)))
624

    
625
def NamedValues(name, cols, values):
626
    '''
627
    @param cols None|[...]
628
    @post `cols` will be changed to Col objects with the table set to `name`.
629
    '''
630
    table = NamedTable(name, Values(values), cols)
631
    if cols != None: set_cols_table(table, cols)
632
    return table
633

    
634
##### Database structure
635

    
636
class TypedCol(Col):
637
    def __init__(self, name, type_, default=None, nullable=True,
638
        constraints=None):
639
        assert default == None or isinstance(default, Code)
640
        
641
        Col.__init__(self, name)
642
        
643
        self.type = type_
644
        self.default = default
645
        self.nullable = nullable
646
        self.constraints = constraints
647
    
648
    def to_str(self, db):
649
        str_ = Col.to_str(self, db)+' '+self.type
650
        if not self.nullable: str_ += ' NOT NULL'
651
        if self.default != None: str_ += ' DEFAULT '+self.default.to_str(db)
652
        if self.constraints != None: str_ += ' '+self.constraints
653
        return str_
654
    
655
    def to_Col(self): return Col(self.name)
656

    
657
ensure_not_null_excs = (NoUnderlyingTableException, KeyError)
658

    
659
def ensure_not_null(db, col, type_=None):
660
    '''
661
    @param col If type_ is not set, must have an underlying column.
662
    @param type_ If set, overrides the underlying column's type.
663
    @return EnsureNotNull|Col
664
    @throws ensure_not_null_excs
665
    '''
666
    nullable = True
667
    try: typed_col = db.col_info(underlying_col(col))
668
    except NoUnderlyingTableException:
669
        if type_ == None: raise
670
    else:
671
        if type_ == None: type_ = typed_col.type
672
        nullable = typed_col.nullable
673
    
674
    if nullable:
675
        try: col = EnsureNotNull(col, type_)
676
        except KeyError, e:
677
            # Warn of no null sentinel for type, even if caller catches error
678
            warnings.warn(UserWarning(exc.str_(e)))
679
            raise
680
    
681
    return col
(25-25/36)