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
    before, sep, version = str_.rpartition('#')
26
    if sep != '': # found match
27
        str_ = before
28
        suffix = sep+version+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
##### General SQL code objects
49

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

    
58
mockDb = MockDb()
59

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

    
65
##### Unparameterized code objects
66

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

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

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

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

    
91
##### Literal values
92

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

    
98
def as_Value(value):
99
    if isinstance(value, Code): return value
100
    else: return Literal(value)
101

    
102
def is_null(value): return isinstance(value, Literal) and value.value == None
103

    
104
##### Derived elements
105

    
106
src_self = object() # tells Col that it is its own source column
107

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

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

    
132
##### Tables
133

    
134
class Table(Derived):
135
    def __init__(self, name, schema=None, srcs=()):
136
        '''
137
        @param schema str|None (for no schema)
138
        @param srcs (Table...)|src_self See Derived.set_srcs()
139
        '''
140
        Derived.__init__(self, srcs)
141
        
142
        name = truncate(name)
143
        
144
        self.name = name
145
        self.schema = schema
146
    
147
    def to_str(self, db):
148
        str_ = ''
149
        if self.schema != None: str_ += db.esc_name(self.schema)+'.'
150
        str_ += db.esc_name(self.name)
151
        return str_
152
    
153
    def to_Table(self): return self
154

    
155
def is_underlying_table(table):
156
    return isinstance(table, Table) and table.to_Table() is table
157

    
158
class NoUnderlyingTableException(Exception): pass
159

    
160
def underlying_table(table):
161
    table = remove_table_rename(table)
162
    if not is_underlying_table(table): raise NoUnderlyingTableException
163
    return table
164

    
165
def as_Table(table, schema=None):
166
    if table == None or isinstance(table, Code): return table
167
    else: return Table(table, schema)
168

    
169
def suffixed_table(table, suffix): return Table(table.name+suffix, table.schema)
170

    
171
class NamedTable(Table):
172
    def __init__(self, name, code, cols=None):
173
        Table.__init__(self, name)
174
        
175
        if not isinstance(code, Code): code = Table(code)
176
        if not isinstance(code, (Table, FunctionCall, Expr)): code = Expr(code)
177
        if cols != None: cols = map(to_name_only_col, cols)
178
        
179
        self.code = code
180
        self.cols = cols
181
    
182
    def to_str(self, db):
183
        str_ = self.code.to_str(db)+'\nAS '+Table.to_str(self, db)
184
        if self.cols != None:
185
            str_ += ' ('+(', '.join((c.to_str(db) for c in self.cols)))+')'
186
        return str_
187
    
188
    def to_Table(self): return Table(self.name)
189

    
190
def remove_table_rename(table):
191
    if isinstance(table, NamedTable): table = table.code
192
    return table
193

    
194
##### Columns
195

    
196
class Col(Derived):
197
    def __init__(self, name, table=None, srcs=()):
198
        '''
199
        @param table Table|None (for no table)
200
        @param srcs (Col...)|src_self See Derived.set_srcs()
201
        '''
202
        Derived.__init__(self, srcs)
203
        
204
        name = truncate(name)
205
        if util.is_str(table): table = Table(table)
206
        assert table == None or isinstance(table, Table)
207
        
208
        self.name = name
209
        self.table = table
210
    
211
    def to_str(self, db, use_concat=False):
212
        str_ = db.esc_name(self.name)
213
        if self.table != None:
214
            str_ = '.'+str_
215
            table_str = self.table.to_Table().to_str(db)
216
            if use_concat: str_ = concat(table_str, str_)
217
            else: str_ = table_str+str_
218
        return str_
219
    
220
    def __str__(self): return clean_name(self.to_str(mockDb, use_concat=True))
221
    
222
    def to_Col(self): return self
223

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

    
226
def as_Col(col, table=None, name=None):
227
    '''
228
    @param name If not None, any non-Col input will be renamed using NamedCol.
229
    '''
230
    if name != None:
231
        col = as_Value(col)
232
        if not isinstance(col, Col): col = NamedCol(name, col)
233
    
234
    if isinstance(col, Code): return col
235
    else: return Col(col, table)
236

    
237
def with_default_table(col, table, overwrite=False):
238
    col = as_Col(col)
239
    if not isinstance(col, NamedCol) and (overwrite or col.table == None):
240
        col = copy.copy(col) # don't modify input!
241
        col.table = table
242
    return col
243

    
244
def set_cols_table(table, cols):
245
    table = as_Table(table)
246
    
247
    for i, col in enumerate(cols):
248
        col = cols[i] = as_Col(col)
249
        col.table = table
250

    
251
def to_name_only_col(col, check_table=None):
252
    col = as_Col(col)
253
    if not isinstance(col, Col): return col
254
    
255
    if check_table != None:
256
        table = col.table
257
        assert table == None or table == check_table
258
    return Col(col.name)
259

    
260
class NamedCol(Col):
261
    def __init__(self, name, code):
262
        Col.__init__(self, name)
263
        
264
        if not isinstance(code, Code): code = Literal(code)
265
        
266
        self.code = code
267
    
268
    def to_str(self, db):
269
        return self.code.to_str(db)+' AS '+Col.to_str(self, db)
270
    
271
    def to_Col(self): return Col(self.name)
272

    
273
def remove_col_rename(col):
274
    if isinstance(col, NamedCol): col = col.code
275
    return col
276

    
277
def underlying_col(col):
278
    col = remove_col_rename(col)
279
    if not isinstance(col, Col): raise NoUnderlyingTableException
280
    
281
    return Col(col.name, underlying_table(col.table), col.srcs)
282

    
283
def wrap(wrap_func, value):
284
    '''Wraps a value, propagating any column renaming to the returned value.'''
285
    if isinstance(value, NamedCol):
286
        return NamedCol(value.name, wrap_func(value.code))
287
    else: return wrap_func(value)
288

    
289
class ColDict(dicts.DictProxy):
290
    '''A dict that automatically makes inserted entries Col objects'''
291
    
292
    def __init__(self, db, keys_table, dict_={}):
293
        dicts.DictProxy.__init__(self, {})
294
        
295
        keys_table = as_Table(keys_table)
296
        
297
        self.db = db
298
        self.table = keys_table
299
        self.update(dict_) # after setting vars because __setitem__() needs them
300
    
301
    def copy(self): return ColDict(self.db, self.table, self.inner.copy())
302
    
303
    def __getitem__(self, key):
304
        return dicts.DictProxy.__getitem__(self, self._key(key))
305
    
306
    def __setitem__(self, key, value):
307
        key = self._key(key)
308
        if value == None: value = self.db.col_info(key).default
309
        dicts.DictProxy.__setitem__(self, key, as_Col(value, name=key.name))
310
    
311
    def _key(self, key): return as_Col(key, self.table)
312

    
313
##### Functions
314

    
315
Function = Table
316
as_Function = as_Table
317

    
318
class InternalFunction(CustomCode): pass
319

    
320
class NamedArg(NamedCol):
321
    def __init__(self, name, value):
322
        NamedCol.__init__(self, name, value)
323
    
324
    def to_str(self, db):
325
        return Col.to_str(self, db)+' := '+self.code.to_str(db)
326

    
327
class FunctionCall(Code):
328
    def __init__(self, function, *args, **kw_args):
329
        '''
330
        @param args [Code|literal-value...] The function's arguments
331
        '''
332
        if not isinstance(function, Code): function = Function(function)
333
        def filter_(arg): return remove_col_rename(as_Value(arg))
334
        args = map(filter_, args)
335
        args += [NamedArg(k, filter_(v)) for k, v in kw_args.iteritems()]
336
        
337
        self.function = function
338
        self.args = args
339
    
340
    def to_str(self, db):
341
        args_str = ', '.join((v.to_str(db) for v in self.args))
342
        return self.function.to_str(db)+'('+args_str+')'
343

    
344
def wrap_in_func(function, value):
345
    '''Wraps a value inside a function call.
346
    Propagates any column renaming to the returned value.
347
    '''
348
    return wrap(lambda v: FunctionCall(function, v), value)
349

    
350
def unwrap_func_call(func_call, check_name=None):
351
    '''Unwraps any function call to its first argument.
352
    Also removes any column renaming.
353
    '''
354
    func_call = remove_col_rename(func_call)
355
    if not isinstance(func_call, FunctionCall): return func_call
356
    
357
    if check_name != None:
358
        name = func_call.function.name
359
        assert name == None or name == check_name
360
    return func_call.args[0]
361

    
362
##### Conditions
363

    
364
class ColValueCond(Code):
365
    def __init__(self, col, value):
366
        value = as_ValueCond(value)
367
        
368
        self.col = col
369
        self.value = value
370
    
371
    def to_str(self, db): return self.value.to_str(db, self.col)
372

    
373
def combine_conds(conds, keyword=None):
374
    '''
375
    @param keyword The keyword to add before the conditions, if any
376
    '''
377
    str_ = ''
378
    if keyword != None:
379
        if conds == []: whitespace = ''
380
        elif len(conds) == 1: whitespace = ' '
381
        else: whitespace = '\n'
382
        str_ += keyword+whitespace
383
    
384
    str_ += '\nAND '.join(conds)
385
    return str_
386

    
387
##### Condition column comparisons
388

    
389
class ValueCond(BasicObject):
390
    def __init__(self, value):
391
        value = remove_col_rename(as_Value(value))
392
        
393
        self.value = value
394
    
395
    def to_str(self, db, left_value):
396
        '''
397
        @param left_value The Code object that the condition is being applied on
398
        '''
399
        raise NotImplemented()
400
    
401
    def __repr__(self): return self.to_str(mockDb, '<left_value>')
402

    
403
class CompareCond(ValueCond):
404
    def __init__(self, value, operator='='):
405
        '''
406
        @param operator By default, compares NULL values literally. Use '~=' or
407
            '~!=' to pass NULLs through.
408
        '''
409
        ValueCond.__init__(self, value)
410
        self.operator = operator
411
    
412
    def to_str(self, db, left_value):
413
        left_value = remove_col_rename(as_Col(left_value))
414
        
415
        right_value = self.value
416
        
417
        # Parse operator
418
        operator = self.operator
419
        passthru_null_ref = [False]
420
        operator = strings.remove_prefix('~', operator, passthru_null_ref)
421
        neg_ref = [False]
422
        operator = strings.remove_prefix('!', operator, neg_ref)
423
        equals = operator.endswith('=') # also includes <=, >=
424
        
425
        # Handle nullable columns
426
        check_null = False
427
        if not passthru_null_ref[0]: # NULLs compare equal
428
            try: left_value = ensure_not_null(db, left_value)
429
            except ensure_not_null_excs: # fall back to alternate method
430
                check_null = equals and isinstance(right_value, Col)
431
            else:
432
                if isinstance(left_value, EnsureNotNull):
433
                    right_value = ensure_not_null(db, right_value,
434
                        left_value.type) # apply same function to both sides
435
        
436
        if equals and is_null(right_value): operator = 'IS'
437
        
438
        left = left_value.to_str(db)
439
        right = right_value.to_str(db)
440
        
441
        # Create str
442
        str_ = left+' '+operator+' '+right
443
        if check_null:
444
            str_ = '('+str_+' OR ('+left+' IS NULL AND '+right+' IS NULL))'
445
        if neg_ref[0]: str_ = 'NOT '+str_
446
        return str_
447

    
448
# Tells as_ValueCond() to assume a non-ValueCond is a literal value
449
assume_literal = object()
450

    
451
def as_ValueCond(value, default_table=assume_literal):
452
    if not isinstance(value, ValueCond):
453
        if default_table is not assume_literal:
454
            value = with_default_table(value, default_table)
455
        return CompareCond(value)
456
    else: return value
457

    
458
##### Joins
459

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

    
462
# Tells Join the left and right columns have the same name and are never NULL
463
join_same_not_null = object()
464

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

    
467
class Join(BasicObject):
468
    def __init__(self, table, mapping={}, type_=None):
469
        '''
470
        @param mapping dict(right_table_col=left_table_col, ...)
471
            * if left_table_col is join_same: left_table_col = right_table_col
472
              * Note that right_table_col must be a string
473
            * if left_table_col is join_same_not_null:
474
              left_table_col = right_table_col and both have NOT NULL constraint
475
              * Note that right_table_col must be a string
476
        @param type_ None (for plain join)|str (e.g. 'LEFT')|filter_out
477
            * filter_out: equivalent to 'LEFT' with the query filtered by
478
              `table_pkey IS NULL` (indicating no match)
479
        '''
480
        if util.is_str(table): table = Table(table)
481
        assert type_ == None or util.is_str(type_) or type_ is filter_out
482
        
483
        self.table = table
484
        self.mapping = mapping
485
        self.type_ = type_
486
    
487
    def to_str(self, db, left_table_):
488
        def join(entry):
489
            '''Parses non-USING joins'''
490
            right_table_col, left_table_col = entry
491
            
492
            # Switch order (right_table_col is on the left in the comparison)
493
            left = right_table_col
494
            right = left_table_col
495
            left_table = self.table
496
            right_table = left_table_
497
            
498
            # Parse left side
499
            left = with_default_table(left, left_table)
500
            
501
            # Parse special values
502
            left_on_right = Col(left.name, right_table)
503
            if right is join_same: right = left_on_right
504
            elif right is join_same_not_null:
505
                right = CompareCond(left_on_right, '~=')
506
            
507
            # Parse right side
508
            right = as_ValueCond(right, right_table)
509
            
510
            return right.to_str(db, left)
511
        
512
        # Create join condition
513
        type_ = self.type_
514
        joins = self.mapping
515
        if joins == {}: join_cond = None
516
        elif type_ is not filter_out and reduce(operator.and_,
517
            (v is join_same_not_null for v in joins.itervalues())):
518
            # all cols w/ USING, so can use simpler USING syntax
519
            cols = map(to_name_only_col, joins.iterkeys())
520
            join_cond = 'USING ('+(', '.join((c.to_str(db) for c in cols)))+')'
521
        else: join_cond = combine_conds(map(join, joins.iteritems()), 'ON')
522
        
523
        if isinstance(self.table, NamedTable): whitespace = '\n'
524
        else: whitespace = ' '
525
        
526
        # Create join
527
        if type_ is filter_out: type_ = 'LEFT'
528
        str_ = ''
529
        if type_ != None: str_ += type_+' '
530
        str_ += 'JOIN'+whitespace+self.table.to_str(db)
531
        if join_cond != None: str_ += whitespace+join_cond
532
        return str_
533
    
534
    def __repr__(self): return self.to_str(mockDb, '<left_table>')
535

    
536
##### Value exprs
537

    
538
default = CustomCode('DEFAULT')
539

    
540
row_count = CustomCode('count(*)')
541

    
542
# See <http://www.postgresql.org/docs/8.3/static/datatype-numeric.html>
543
null_sentinels = {
544
    'character varying': r'\N',
545
    'double precision': 'NaN',
546
    'integer': 2147483647,
547
    'text': r'\N',
548
    'timestamp with time zone': 'infinity'
549
}
550

    
551
class EnsureNotNull(FunctionCall):
552
    def __init__(self, value, type_):
553
        FunctionCall.__init__(self, InternalFunction('COALESCE'), as_Col(value),
554
            null_sentinels[type_])
555
        
556
        self.type = type_
557

    
558
##### Table exprs
559

    
560
class Values(Code):
561
    def __init__(self, values):
562
        '''
563
        @param values [...]|[[...], ...] Can be one or multiple rows.
564
        '''
565
        rows = values
566
        if len(values) >= 1 and not lists.is_seq(values[0]): # only one row
567
            rows = [values]
568
        for i, row in enumerate(rows):
569
            rows[i] = map(remove_col_rename, map(as_Value, row))
570
        
571
        self.rows = rows
572
    
573
    def to_str(self, db):
574
        def row_str(row):
575
            return '('+(', '.join((v.to_str(db) for v in row)))+')'
576
        return 'VALUES '+(', '.join(map(row_str, self.rows)))
577

    
578
def NamedValues(name, cols, values):
579
    '''
580
    @post `cols` will be changed to Col objects with the table set to `name`.
581
    '''
582
    table = NamedTable(name, Values(values), cols)
583
    set_cols_table(table, cols)
584
    return table
585

    
586
##### Database structure
587

    
588
class TypedCol(Col):
589
    def __init__(self, name, type_, default=None, nullable=True,
590
        constraints=None):
591
        assert default == None or isinstance(default, Code)
592
        
593
        Col.__init__(self, name)
594
        
595
        self.type = type_
596
        self.default = default
597
        self.nullable = nullable
598
        self.constraints = constraints
599
    
600
    def to_str(self, db):
601
        str_ = Col.to_str(self, db)+' '+self.type
602
        if not self.nullable: str_ += ' NOT NULL'
603
        if self.default != None: str_ += ' DEFAULT '+self.default.to_str(db)
604
        if self.constraints != None: str_ += ' '+self.constraints
605
        return str_
606
    
607
    def to_Col(self): return Col(self.name)
608

    
609
ensure_not_null_excs = (NoUnderlyingTableException, KeyError)
610

    
611
def ensure_not_null(db, col, type_=None):
612
    '''
613
    @param col If type_ is not set, must have an underlying column.
614
    @param type_ If set, overrides the underlying column's type.
615
    @return EnsureNotNull|Col
616
    @throws ensure_not_null_excs
617
    '''
618
    nullable = True
619
    try: typed_col = db.col_info(underlying_col(col))
620
    except NoUnderlyingTableException:
621
        if type_ == None: raise
622
    else:
623
        if type_ == None: type_ = typed_col.type
624
        nullable = typed_col.nullable
625
    
626
    if nullable:
627
        try: col = EnsureNotNull(col, type_)
628
        except KeyError, e:
629
            # Warn of no null sentinel for type, even if caller catches error
630
            warnings.warn(UserWarning(exc.str_(e)))
631
            raise
632
    
633
    return col
(25-25/36)