Project

General

Profile

1
# Database access
2

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

    
8
import exc
9
import dicts
10
import iters
11
import lists
12
from Proxy import Proxy
13
import rand
14
import sql_gen
15
import strings
16
import util
17

    
18
##### Exceptions
19

    
20
def get_cur_query(cur, input_query=None, input_params=None):
21
    raw_query = None
22
    if hasattr(cur, 'query'): raw_query = cur.query
23
    elif hasattr(cur, '_last_executed'): raw_query = cur._last_executed
24
    
25
    if raw_query != None: return raw_query
26
    else: return '[input] '+strings.ustr(input_query)+' % '+repr(input_params)
27

    
28
def _add_cursor_info(e, *args, **kw_args):
29
    '''For params, see get_cur_query()'''
30
    exc.add_msg(e, 'query: '+str(get_cur_query(*args, **kw_args)))
31

    
32
class DbException(exc.ExceptionWithCause):
33
    def __init__(self, msg, cause=None, cur=None):
34
        exc.ExceptionWithCause.__init__(self, msg, cause, cause_newline=True)
35
        if cur != None: _add_cursor_info(self, cur)
36

    
37
class ExceptionWithName(DbException):
38
    def __init__(self, name, cause=None):
39
        DbException.__init__(self, 'for name: '+strings.as_tt(str(name)), cause)
40
        self.name = name
41

    
42
class ExceptionWithNameValue(DbException):
43
    def __init__(self, name, value, cause=None):
44
        DbException.__init__(self, 'for name: '+strings.as_tt(str(name))
45
            +'; value: '+strings.as_tt(repr(value)), cause)
46
        self.name = name
47
        self.value = value
48

    
49
class ConstraintException(DbException):
50
    def __init__(self, name, cols, cause=None):
51
        DbException.__init__(self, 'Violated '+strings.as_tt(name)
52
            +' constraint on columns: '+strings.as_tt(', '.join(cols)), cause)
53
        self.name = name
54
        self.cols = cols
55

    
56
class MissingCastException(DbException):
57
    def __init__(self, type_, col, cause=None):
58
        DbException.__init__(self, 'Missing cast to type '+strings.as_tt(type_)
59
            +' on column: '+strings.as_tt(col), cause)
60
        self.type = type_
61
        self.col = col
62

    
63
class NameException(DbException): pass
64

    
65
class DuplicateKeyException(ConstraintException): pass
66

    
67
class NullValueException(ConstraintException): pass
68

    
69
class FunctionValueException(ExceptionWithNameValue): pass
70

    
71
class DuplicateTableException(ExceptionWithName): pass
72

    
73
class DuplicateFunctionException(ExceptionWithName): pass
74

    
75
class EmptyRowException(DbException): pass
76

    
77
##### Warnings
78

    
79
class DbWarning(UserWarning): pass
80

    
81
##### Result retrieval
82

    
83
def col_names(cur): return (col[0] for col in cur.description)
84

    
85
def rows(cur): return iter(lambda: cur.fetchone(), None)
86

    
87
def consume_rows(cur):
88
    '''Used to fetch all rows so result will be cached'''
89
    iters.consume_iter(rows(cur))
90

    
91
def next_row(cur): return rows(cur).next()
92

    
93
def row(cur):
94
    row_ = next_row(cur)
95
    consume_rows(cur)
96
    return row_
97

    
98
def next_value(cur): return next_row(cur)[0]
99

    
100
def value(cur): return row(cur)[0]
101

    
102
def values(cur): return iters.func_iter(lambda: next_value(cur))
103

    
104
def value_or_none(cur):
105
    try: return value(cur)
106
    except StopIteration: return None
107

    
108
##### Input validation
109

    
110
def esc_name_by_module(module, name):
111
    if module == 'psycopg2' or module == None: quote = '"'
112
    elif module == 'MySQLdb': quote = '`'
113
    else: raise NotImplementedError("Can't escape name for "+module+' database')
114
    return sql_gen.esc_name(name, quote)
115

    
116
def esc_name_by_engine(engine, name, **kw_args):
117
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
118

    
119
def esc_name(db, name, **kw_args):
120
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
121

    
122
def qual_name(db, schema, table):
123
    def esc_name_(name): return esc_name(db, name)
124
    table = esc_name_(table)
125
    if schema != None: return esc_name_(schema)+'.'+table
126
    else: return table
127

    
128
##### Database connections
129

    
130
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
131

    
132
db_engines = {
133
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
134
    'PostgreSQL': ('psycopg2', {}),
135
}
136

    
137
DatabaseErrors_set = set([DbException])
138
DatabaseErrors = tuple(DatabaseErrors_set)
139

    
140
def _add_module(module):
141
    DatabaseErrors_set.add(module.DatabaseError)
142
    global DatabaseErrors
143
    DatabaseErrors = tuple(DatabaseErrors_set)
144

    
145
def db_config_str(db_config):
146
    return db_config['engine']+' database '+db_config['database']
147

    
148
def _query_lookup(query, params): return (query, dicts.make_hashable(params))
149

    
150
log_debug_none = lambda msg, level=2: None
151

    
152
class DbConn:
153
    def __init__(self, db_config, serializable=True, autocommit=False,
154
        caching=True, log_debug=log_debug_none):
155
        self.db_config = db_config
156
        self.serializable = serializable
157
        self.autocommit = autocommit
158
        self.caching = caching
159
        self.log_debug = log_debug
160
        self.debug = log_debug != log_debug_none
161
        
162
        self.__db = None
163
        self.query_results = {}
164
        self._savepoint = 0
165
    
166
    def __getattr__(self, name):
167
        if name == '__dict__': raise Exception('getting __dict__')
168
        if name == 'db': return self._db()
169
        else: raise AttributeError()
170
    
171
    def __getstate__(self):
172
        state = copy.copy(self.__dict__) # shallow copy
173
        state['log_debug'] = None # don't pickle the debug callback
174
        state['_DbConn__db'] = None # don't pickle the connection
175
        return state
176
    
177
    def connected(self): return self.__db != None
178
    
179
    def _db(self):
180
        if self.__db == None:
181
            # Process db_config
182
            db_config = self.db_config.copy() # don't modify input!
183
            schemas = db_config.pop('schemas', None)
184
            module_name, mappings = db_engines[db_config.pop('engine')]
185
            module = __import__(module_name)
186
            _add_module(module)
187
            for orig, new in mappings.iteritems():
188
                try: util.rename_key(db_config, orig, new)
189
                except KeyError: pass
190
            
191
            # Connect
192
            self.__db = module.connect(**db_config)
193
            
194
            # Configure connection
195
            if self.serializable and not self.autocommit: run_raw_query(self,
196
                'SET TRANSACTION ISOLATION LEVEL SERIALIZABLE')
197
            if schemas != None:
198
                schemas_ = ''.join((esc_name(self, s)+', '
199
                    for s in schemas.split(',')))
200
                run_raw_query(self, "SELECT set_config('search_path', \
201
%s || current_setting('search_path'), false)", [schemas_])
202
        
203
        return self.__db
204
    
205
    class DbCursor(Proxy):
206
        def __init__(self, outer):
207
            Proxy.__init__(self, outer.db.cursor())
208
            self.outer = outer
209
            self.query_results = outer.query_results
210
            self.query_lookup = None
211
            self.result = []
212
        
213
        def execute(self, query, params=None):
214
            self._is_insert = query.upper().find('INSERT') >= 0
215
            self.query_lookup = _query_lookup(query, params)
216
            try:
217
                try:
218
                    return_value = self.inner.execute(query, params)
219
                    self.outer.do_autocommit()
220
                finally: self.query = get_cur_query(self.inner, query, params)
221
            except Exception, e:
222
                _add_cursor_info(e, self, query, params)
223
                self.result = e # cache the exception as the result
224
                self._cache_result()
225
                raise
226
            # Fetch all rows so result will be cached
227
            if self.rowcount == 0 and not self._is_insert: consume_rows(self)
228
            return return_value
229
        
230
        def fetchone(self):
231
            row = self.inner.fetchone()
232
            if row != None: self.result.append(row)
233
            # otherwise, fetched all rows
234
            else: self._cache_result()
235
            return row
236
        
237
        def _cache_result(self):
238
            # For inserts, only cache exceptions since inserts are not
239
            # idempotent, but an invalid insert will always be invalid
240
            if self.query_results != None and (not self._is_insert
241
                or isinstance(self.result, Exception)):
242
                
243
                assert self.query_lookup != None
244
                self.query_results[self.query_lookup] = self.CacheCursor(
245
                    util.dict_subset(dicts.AttrsDictView(self),
246
                    ['query', 'result', 'rowcount', 'description']))
247
        
248
        class CacheCursor:
249
            def __init__(self, cached_result): self.__dict__ = cached_result
250
            
251
            def execute(self, *args, **kw_args):
252
                if isinstance(self.result, Exception): raise self.result
253
                # otherwise, result is a rows list
254
                self.iter = iter(self.result)
255
            
256
            def fetchone(self):
257
                try: return self.iter.next()
258
                except StopIteration: return None
259
    
260
    def esc_value(self, value):
261
        module = util.root_module(self.db)
262
        if module == 'psycopg2': str_ = self.db.cursor().mogrify('%s', [value])
263
        elif module == 'MySQLdb':
264
            import _mysql
265
            str_ = _mysql.escape_string(value)
266
        else: raise NotImplementedError("Can't escape value for "+module
267
            +' database')
268
        return strings.to_unicode(str_)
269
    
270
    def esc_name(self, name): return esc_name(self, name) # calls global func
271
    
272
    def run_query(self, query, params=None, cacheable=False, log_level=2,
273
        debug_msg_ref=None):
274
        '''
275
        @param log_ignore_excs The log_level will be increased by 2 if the query
276
            throws one of these exceptions.
277
        '''
278
        assert query != None
279
        
280
        if not self.caching: cacheable = False
281
        used_cache = False
282
        try:
283
            # Get cursor
284
            if cacheable:
285
                query_lookup = _query_lookup(query, params)
286
                try:
287
                    cur = self.query_results[query_lookup]
288
                    used_cache = True
289
                except KeyError: cur = self.DbCursor(self)
290
            else: cur = self.db.cursor()
291
            
292
            # Run query
293
            cur.execute(query, params)
294
        finally:
295
            if self.debug and debug_msg_ref != None:# only compute msg if needed
296
                if used_cache: cache_status = 'cache hit'
297
                elif cacheable: cache_status = 'cache miss'
298
                else: cache_status = 'non-cacheable'
299
                query_code = strings.as_code(str(get_cur_query(cur, query,
300
                    params)), 'SQL')
301
                debug_msg_ref[0] = 'DB query: '+cache_status+':\n'+query_code
302
        
303
        return cur
304
    
305
    def is_cached(self, query, params=None):
306
        return _query_lookup(query, params) in self.query_results
307
    
308
    def with_savepoint(self, func):
309
        savepoint = 'level_'+str(self._savepoint)
310
        self.run_query('SAVEPOINT '+savepoint, log_level=4)
311
        self._savepoint += 1
312
        try: 
313
            try: return_val = func()
314
            finally:
315
                self._savepoint -= 1
316
                assert self._savepoint >= 0
317
        except:
318
            self.run_query('ROLLBACK TO SAVEPOINT '+savepoint, log_level=4)
319
            raise
320
        else:
321
            self.run_query('RELEASE SAVEPOINT '+savepoint, log_level=4)
322
            self.do_autocommit()
323
            return return_val
324
    
325
    def do_autocommit(self):
326
        '''Autocommits if outside savepoint'''
327
        assert self._savepoint >= 0
328
        if self.autocommit and self._savepoint == 0:
329
            self.log_debug('Autocommiting')
330
            self.db.commit()
331
    
332
    def col_default(self, col):
333
        table = sql_gen.Table('columns', 'information_schema')
334
        
335
        conds = [('table_name', col.table.name), ('column_name', col.name)]
336
        schema = col.table.schema
337
        if schema != None: conds.append(('table_schema', schema))
338
        
339
        return value(select(self, table, ['column_default'], conds,
340
            order_by='table_schema', limit=1, log_level=3))
341
            # TODO: order_by search_path schema order
342

    
343
connect = DbConn
344

    
345
##### Querying
346

    
347
def run_raw_query(db, *args, **kw_args):
348
    '''For params, see DbConn.run_query()'''
349
    return db.run_query(*args, **kw_args)
350

    
351
def mogrify(db, query, params):
352
    module = util.root_module(db.db)
353
    if module == 'psycopg2': return db.db.cursor().mogrify(query, params)
354
    else: raise NotImplementedError("Can't mogrify query for "+module+
355
        ' database')
356

    
357
##### Recoverable querying
358

    
359
def with_savepoint(db, func): return db.with_savepoint(func)
360

    
361
def run_query(db, query, params=None, recover=None, cacheable=False,
362
    log_level=2, log_ignore_excs=None, **kw_args):
363
    '''For params, see run_raw_query()'''
364
    if recover == None: recover = False
365
    if log_ignore_excs == None: log_ignore_excs = ()
366
    log_ignore_excs = tuple(log_ignore_excs)
367
    
368
    debug_msg_ref = [None]
369
    try:
370
        try:
371
            def run(): return run_raw_query(db, query, params, cacheable,
372
                log_level, debug_msg_ref, **kw_args)
373
            if recover and not db.is_cached(query, params):
374
                return with_savepoint(db, run)
375
            else: return run() # don't need savepoint if cached
376
        except Exception, e:
377
            if not recover: raise # need savepoint to run index_cols()
378
            msg = exc.str_(e)
379
            
380
            match = re.search(r'duplicate key value violates unique constraint '
381
                r'"((_?[^\W_]+)_.+?)"', msg)
382
            if match:
383
                constraint, table = match.groups()
384
                try: cols = index_cols(db, table, constraint)
385
                except NotImplementedError: raise e
386
                else: raise DuplicateKeyException(constraint, cols, e)
387
            
388
            match = re.search(r'null value in column "(.+?)" violates not-null'
389
                r' constraint', msg)
390
            if match: raise NullValueException('NOT NULL', [match.group(1)], e)
391
            
392
            match = re.search(r'\b(?:invalid input (?:syntax|value)\b.*?'
393
                r'|date/time field value out of range): "(.+?)"\n'
394
                r'(?:(?s).*?)\bfunction "(.+?)"', msg)
395
            if match:
396
                value, name = match.groups()
397
                raise FunctionValueException(name, strings.to_unicode(value), e)
398
            
399
            match = re.search(r'column "(.+?)" is of type (.+?) but expression '
400
                r'is of type', msg)
401
            if match:
402
                col, type_ = match.groups()
403
                raise MissingCastException(type_, col, e)
404
            
405
            match = re.search(r'relation "(.+?)" already exists', msg)
406
            if match: raise DuplicateTableException(match.group(1), e)
407
            
408
            match = re.search(r'function "(.+?)" already exists', msg)
409
            if match: raise DuplicateFunctionException(match.group(1), e)
410
            
411
            raise # no specific exception raised
412
    except log_ignore_excs:
413
        log_level += 2
414
        raise
415
    finally:
416
        if debug_msg_ref[0] != None: db.log_debug(debug_msg_ref[0], log_level)
417

    
418
##### Basic queries
419

    
420
def next_version(name):
421
    version = 1 # first existing name was version 0
422
    match = re.match(r'^(.*)#(\d+)$', name)
423
    if match:
424
        name, version = match.groups()
425
        version = int(version)+1
426
    return sql_gen.add_suffix(name, '#'+str(version))
427

    
428
def run_query_into(db, query, params, into=None, *args, **kw_args):
429
    '''Outputs a query to a temp table.
430
    For params, see run_query().
431
    '''
432
    if into == None: return run_query(db, query, params, *args, **kw_args)
433
    else: # place rows in temp table
434
        assert isinstance(into, sql_gen.Table)
435
        
436
        kw_args['recover'] = True
437
        kw_args.setdefault('log_ignore_excs', (DuplicateTableException,))
438
        
439
        temp = not db.autocommit # tables are permanent in autocommit mode
440
        # "temporary tables cannot specify a schema name", so remove schema
441
        if temp: into.schema = None
442
        
443
        while True:
444
            try:
445
                create_query = 'CREATE'
446
                if temp: create_query += ' TEMP'
447
                create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
448
                
449
                return run_query(db, create_query, params, *args, **kw_args)
450
                    # CREATE TABLE AS sets rowcount to # rows in query
451
            except DuplicateTableException, e:
452
                into.name = next_version(into.name)
453
                # try again with next version of name
454

    
455
order_by_pkey = object() # tells mk_select() to order by the pkey
456

    
457
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
458

    
459
def mk_select(db, tables, fields=None, conds=None, distinct_on=[], limit=None,
460
    start=None, order_by=order_by_pkey, default_table=None):
461
    '''
462
    @param tables The single table to select from, or a list of tables to join
463
        together, with tables after the first being sql_gen.Join objects
464
    @param fields Use None to select all fields in the table
465
    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
466
        * container can be any iterable type
467
        * compare_left_side: sql_gen.Code|str (for col name)
468
        * compare_right_side: sql_gen.ValueCond|literal value
469
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
470
        use all columns
471
    @return tuple(query, params)
472
    '''
473
    # Parse tables param
474
    if not lists.is_seq(tables): tables = [tables]
475
    tables = list(tables) # don't modify input! (list() copies input)
476
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
477
    
478
    # Parse other params
479
    if conds == None: conds = []
480
    elif dicts.is_dict(conds): conds = conds.items()
481
    conds = list(conds) # don't modify input! (list() copies input)
482
    assert limit == None or type(limit) == int
483
    assert start == None or type(start) == int
484
    if order_by is order_by_pkey:
485
        if distinct_on != []: order_by = None
486
        else: order_by = pkey(db, table0, recover=True)
487
    
488
    query = 'SELECT'
489
    
490
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
491
    
492
    # DISTINCT ON columns
493
    if distinct_on != []:
494
        query += '\nDISTINCT'
495
        if distinct_on is not distinct_on_all:
496
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
497
    
498
    # Columns
499
    query += '\n'
500
    if fields == None: query += '*'
501
    else: query += '\n, '.join(map(parse_col, fields))
502
    
503
    # Main table
504
    query += '\nFROM '+table0.to_str(db)
505
    
506
    # Add joins
507
    left_table = table0
508
    for join_ in tables:
509
        table = join_.table
510
        
511
        # Parse special values
512
        if join_.type_ is sql_gen.filter_out: # filter no match
513
            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
514
                None))
515
        
516
        query += '\n'+join_.to_str(db, left_table)
517
        
518
        left_table = table
519
    
520
    missing = True
521
    if conds != []:
522
        if len(conds) == 1: whitespace = ' '
523
        else: whitespace = '\n'
524
        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
525
            .to_str(db) for l, r in conds], 'WHERE')
526
        missing = False
527
    if order_by != None:
528
        query += '\nORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)
529
    if limit != None: query += '\nLIMIT '+str(limit); missing = False
530
    if start != None:
531
        if start != 0: query += '\nOFFSET '+str(start)
532
        missing = False
533
    if missing: warnings.warn(DbWarning(
534
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
535
    
536
    return (query, [])
537

    
538
def select(db, *args, **kw_args):
539
    '''For params, see mk_select() and run_query()'''
540
    recover = kw_args.pop('recover', None)
541
    cacheable = kw_args.pop('cacheable', True)
542
    log_level = kw_args.pop('log_level', 2)
543
    
544
    query, params = mk_select(db, *args, **kw_args)
545
    return run_query(db, query, params, recover, cacheable, log_level=log_level)
546

    
547
def mk_insert_select(db, table, cols=None, select_query=None, params=None,
548
    returning=None, embeddable=False):
549
    '''
550
    @param returning str|None An inserted column (such as pkey) to return
551
    @param embeddable Whether the query should be embeddable as a nested SELECT.
552
        Warning: If you set this and cacheable=True when the query is run, the
553
        query will be fully cached, not just if it raises an exception.
554
    '''
555
    table = sql_gen.as_Table(table)
556
    if cols == []: cols = None # no cols (all defaults) = unknown col names
557
    if cols != None:
558
        cols = [sql_gen.to_name_only_col(v, table).to_str(db) for v in cols]
559
    if select_query == None: select_query = 'DEFAULT VALUES'
560
    if returning != None: returning = sql_gen.as_Col(returning, table)
561
    
562
    # Build query
563
    first_line = 'INSERT INTO '+table.to_str(db)
564
    query = first_line
565
    if cols != None: query += '\n('+', '.join(cols)+')'
566
    query += '\n'+select_query
567
    
568
    if returning != None:
569
        returning_name = copy.copy(returning)
570
        returning_name.table = None
571
        returning_name = returning_name.to_str(db)
572
        query += '\nRETURNING '+returning_name
573
    
574
    if embeddable:
575
        assert returning != None
576
        
577
        # Create function
578
        function_name = sql_gen.clean_name(first_line)
579
        return_type = 'SETOF '+returning.to_str(db)+'%TYPE'
580
        while True:
581
            try:
582
                func_schema = None
583
                if not db.autocommit: func_schema = 'pg_temp'
584
                function = sql_gen.Table(function_name, func_schema).to_str(db)
585
                
586
                function_query = '''\
587
CREATE FUNCTION '''+function+'''()
588
RETURNS '''+return_type+'''
589
LANGUAGE sql
590
AS $$
591
'''+mogrify(db, query, params)+''';
592
$$;
593
'''
594
                run_query(db, function_query, recover=True, cacheable=True,
595
                    log_ignore_excs=(DuplicateFunctionException,))
596
                break # this version was successful
597
            except DuplicateFunctionException, e:
598
                function_name = next_version(function_name)
599
                # try again with next version of name
600
        
601
        # Return query that uses function
602
        func_table = sql_gen.NamedTable('f', sql_gen.CustomCode(function+'()'),
603
            [returning_name]) # AS clause requires function alias
604
        return mk_select(db, func_table, start=0, order_by=None)
605
    
606
    return (query, params)
607

    
608
def insert_select(db, *args, **kw_args):
609
    '''For params, see mk_insert_select() and run_query_into()
610
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
611
        values in
612
    '''
613
    into = kw_args.pop('into', None)
614
    if into != None: kw_args['embeddable'] = True
615
    recover = kw_args.pop('recover', None)
616
    cacheable = kw_args.pop('cacheable', True)
617
    
618
    query, params = mk_insert_select(db, *args, **kw_args)
619
    return run_query_into(db, query, params, into, recover=recover,
620
        cacheable=cacheable)
621

    
622
default = object() # tells insert() to use the default value for a column
623

    
624
def insert(db, table, row, *args, **kw_args):
625
    '''For params, see insert_select()'''
626
    if lists.is_seq(row): cols = None
627
    else:
628
        cols = row.keys()
629
        row = row.values()
630
    row = list(row) # ensure that "!= []" works
631
    
632
    # Check for special values
633
    labels = []
634
    values = []
635
    for value in row:
636
        value = sql_gen.remove_col_rename(sql_gen.as_Value(value)).value
637
        if value is default: labels.append('DEFAULT')
638
        else:
639
            labels.append('%s')
640
            values.append(value)
641
    
642
    # Build query
643
    if values != []: query = 'VALUES ('+(', '.join(labels))+')'
644
    else: query = None
645
    
646
    return insert_select(db, table, cols, query, values, *args, **kw_args)
647

    
648
def mk_update(db, table, changes=None, cond=None):
649
    '''
650
    @param changes [(col, new_value),...]
651
        * container can be any iterable type
652
        * col: sql_gen.Code|str (for col name)
653
        * new_value: sql_gen.Code|literal value
654
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
655
    @return str query
656
    '''
657
    query = 'UPDATE '+sql_gen.as_Table(table).to_str(db)+'\nSET\n'
658
    query += ',\n'.join((sql_gen.to_name_only_col(col, table).to_str(db)+' = '
659
        +sql_gen.as_Value(new_value).to_str(db) for col, new_value in changes))
660
    if cond != None: query += '\nWHERE\n'+cond.to_str(db)
661
    
662
    return query
663

    
664
def update(db, *args, **kw_args):
665
    '''For params, see mk_update() and run_query()'''
666
    recover = kw_args.pop('recover', None)
667
    
668
    return run_query(db, mk_update(db, *args, **kw_args), [], recover)
669

    
670
def last_insert_id(db):
671
    module = util.root_module(db.db)
672
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
673
    elif module == 'MySQLdb': return db.insert_id()
674
    else: return None
675

    
676
def truncate(db, table, schema='public'):
677
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
678

    
679
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
680
    '''Creates a mapping from original column names (which may have collisions)
681
    to names that will be distinct among the columns' tables.
682
    This is meant to be used for several tables that are being joined together.
683
    @param cols The columns to combine. Duplicates will be removed.
684
    @param into The table for the new columns.
685
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
686
        columns will be included in the mapping even if they are not in cols.
687
        The tables of the provided Col objects will be changed to into, so make
688
        copies of them if you want to keep the original tables.
689
    @param as_items Whether to return a list of dict items instead of a dict
690
    @return dict(orig_col=new_col, ...)
691
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
692
        * new_col: sql_gen.Col(orig_col_name, into)
693
        * All mappings use the into table so its name can easily be
694
          changed for all columns at once
695
    '''
696
    cols = lists.uniqify(cols)
697
    
698
    items = []
699
    for col in preserve:
700
        orig_col = copy.copy(col)
701
        col.table = into
702
        items.append((orig_col, col))
703
    preserve = set(preserve)
704
    for col in cols:
705
        if col not in preserve: items.append((col, sql_gen.Col(str(col), into)))
706
    
707
    if not as_items: items = dict(items)
708
    return items
709

    
710
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
711
    '''For params, see mk_flatten_mapping()
712
    @return See return value of mk_flatten_mapping()
713
    '''
714
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
715
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
716
    run_query_into(db, *mk_select(db, joins, cols, limit=limit, start=start),
717
        into=into)
718
    return dict(items)
719

    
720
##### Database structure queries
721

    
722
def table_row_count(db, table, recover=None):
723
    return value(run_query(db, *mk_select(db, table, [sql_gen.row_count],
724
        order_by=None, start=0), recover=recover, log_level=3))
725

    
726
def table_cols(db, table, recover=None):
727
    return list(col_names(select(db, table, limit=0, order_by=None,
728
        recover=recover, log_level=4)))
729

    
730
def pkey(db, table, recover=None):
731
    '''Assumed to be first column in table'''
732
    return table_cols(db, table, recover)[0]
733

    
734
not_null_col = 'not_null_col'
735

    
736
def table_not_null_col(db, table, recover=None):
737
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
738
    if not_null_col in table_cols(db, table, recover): return not_null_col
739
    else: return pkey(db, table, recover)
740

    
741
def index_cols(db, table, index):
742
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
743
    automatically created. When you don't know whether something is a UNIQUE
744
    constraint or a UNIQUE index, use this function.'''
745
    module = util.root_module(db.db)
746
    if module == 'psycopg2':
747
        return list(values(run_query(db, '''\
748
SELECT attname
749
FROM
750
(
751
        SELECT attnum, attname
752
        FROM pg_index
753
        JOIN pg_class index ON index.oid = indexrelid
754
        JOIN pg_class table_ ON table_.oid = indrelid
755
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
756
        WHERE
757
            table_.relname = %(table)s
758
            AND index.relname = %(index)s
759
    UNION
760
        SELECT attnum, attname
761
        FROM
762
        (
763
            SELECT
764
                indrelid
765
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
766
                    AS indkey
767
            FROM pg_index
768
            JOIN pg_class index ON index.oid = indexrelid
769
            JOIN pg_class table_ ON table_.oid = indrelid
770
            WHERE
771
                table_.relname = %(table)s
772
                AND index.relname = %(index)s
773
        ) s
774
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
775
) s
776
ORDER BY attnum
777
''',
778
            {'table': table, 'index': index}, cacheable=True, log_level=4)))
779
    else: raise NotImplementedError("Can't list index columns for "+module+
780
        ' database')
781

    
782
def constraint_cols(db, table, constraint):
783
    module = util.root_module(db.db)
784
    if module == 'psycopg2':
785
        return list(values(run_query(db, '''\
786
SELECT attname
787
FROM pg_constraint
788
JOIN pg_class ON pg_class.oid = conrelid
789
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
790
WHERE
791
    relname = %(table)s
792
    AND conname = %(constraint)s
793
ORDER BY attnum
794
''',
795
            {'table': table, 'constraint': constraint})))
796
    else: raise NotImplementedError("Can't list constraint columns for "+module+
797
        ' database')
798

    
799
row_num_col = '_row_num'
800

    
801
def add_index(db, expr):
802
    '''Adds an index on a column or expression if it doesn't already exist.
803
    Currently, only function calls are supported as expressions.
804
    '''
805
    expr = copy.copy(expr) # don't modify input!
806
    
807
    # Extract col
808
    if isinstance(expr, sql_gen.FunctionCall):
809
        col = expr.args[0]
810
        expr = sql_gen.Expr(expr)
811
    else: col = expr
812
    assert sql_gen.is_table_col(col)
813
    
814
    index = sql_gen.as_Table(str(expr))
815
    table = col.table
816
    col.table = None
817
    try: run_query(db, 'CREATE INDEX '+index.to_str(db)+' ON '+table.to_str(db)
818
        +' ('+expr.to_str(db)+')', recover=True, cacheable=True, log_level=3)
819
    except DuplicateTableException: pass # index already existed
820

    
821
def add_pkey(db, table, recover=None):
822
    '''Makes the first column in a table the primary key.
823
    @pre The table must not already have a primary key.
824
    '''
825
    table = sql_gen.as_Table(table)
826
    
827
    index = sql_gen.as_Table(sql_gen.add_suffix(table.name, '_pkey'))
828
    col = sql_gen.to_name_only_col(pkey(db, table, recover))
829
    try:
830
        run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD CONSTRAINT '
831
            +index.to_str(db)+' PRIMARY KEY('+col.to_str(db)+')',
832
            recover=True, cacheable=True, log_level=3,
833
            log_ignore_excs=(DuplicateTableException,))
834
    except DuplicateTableException, e:
835
        index.name = next_version(index.name)
836
        # try again with next version of name
837

    
838
def add_row_num(db, table):
839
    '''Adds a row number column to a table. Its name is in row_num_col. It will
840
    be the primary key.'''
841
    table = sql_gen.as_Table(table).to_str(db)
842
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
843
        +' serial NOT NULL PRIMARY KEY', log_level=3)
844

    
845
def tables(db, schema_like='public', table_like='%'):
846
    module = util.root_module(db.db)
847
    params = {'schema_like': schema_like, 'table_like': table_like}
848
    if module == 'psycopg2':
849
        return values(run_query(db, '''\
850
SELECT tablename
851
FROM pg_tables
852
WHERE
853
    schemaname LIKE %(schema_like)s
854
    AND tablename LIKE %(table_like)s
855
ORDER BY tablename
856
''',
857
            params, cacheable=True))
858
    elif module == 'MySQLdb':
859
        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,
860
            cacheable=True))
861
    else: raise NotImplementedError("Can't list tables for "+module+' database')
862

    
863
##### Database management
864

    
865
def empty_db(db, schema='public', **kw_args):
866
    '''For kw_args, see tables()'''
867
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
868

    
869
##### Heuristic queries
870

    
871
def put(db, table, row, pkey_=None, row_ct_ref=None):
872
    '''Recovers from errors.
873
    Only works under PostgreSQL (uses INSERT RETURNING).
874
    '''
875
    row = sql_gen.ColDict(db, table, row)
876
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
877
    
878
    try:
879
        cur = insert(db, table, row, pkey_, recover=True)
880
        if row_ct_ref != None and cur.rowcount >= 0:
881
            row_ct_ref[0] += cur.rowcount
882
        return value(cur)
883
    except DuplicateKeyException, e:
884
        return value(select(db, table, [pkey_],
885
            util.dict_subset_right_join(row, e.cols), recover=True))
886

    
887
def get(db, table, row, pkey, row_ct_ref=None, create=False):
888
    '''Recovers from errors'''
889
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
890
    except StopIteration:
891
        if not create: raise
892
        return put(db, table, row, pkey, row_ct_ref) # insert new row
893

    
894
def is_func_result(col):
895
    return col.table.name.find('(') >= 0 and col.name == 'result'
896

    
897
def into_table_name(out_table, in_tables0, mapping, is_func):
898
    def in_col_str(in_col):
899
        in_col = sql_gen.remove_col_rename(in_col)
900
        if isinstance(in_col, sql_gen.Col):
901
            table = in_col.table
902
            if table == in_tables0:
903
                in_col = sql_gen.to_name_only_col(in_col)
904
            elif is_func_result(in_col): in_col = table # omit col name
905
        return str(in_col)
906
    
907
    str_ = str(out_table)
908
    if is_func:
909
        str_ += '('
910
        
911
        try: value_in_col = mapping['value']
912
        except KeyError:
913
            str_ += ', '.join((k+'='+in_col_str(v)
914
                for k, v in mapping.iteritems()))
915
        else: str_ += in_col_str(value_in_col)
916
        
917
        str_ += ')'
918
    else: str_ += '_pkeys'
919
    return str_
920

    
921
def put_table(db, out_table, in_tables, mapping, row_ct_ref=None, into=None,
922
    default=None, is_func=False):
923
    '''Recovers from errors.
924
    Only works under PostgreSQL (uses INSERT RETURNING).
925
    @param in_tables The main input table to select from, followed by a list of
926
        tables to join with it using the main input table's pkey
927
    @param mapping dict(out_table_col=in_table_col, ...)
928
        * out_table_col: str (*not* sql_gen.Col)
929
        * in_table_col: sql_gen.Col|literal-value
930
    @param into The table to contain the output and input pkeys.
931
        Defaults to `out_table.name+'_pkeys'`.
932
    @param default The *output* column to use as the pkey for missing rows.
933
        If this output column does not exist in the mapping, uses None.
934
    @param is_func Whether out_table is the name of a SQL function, not a table
935
    @return sql_gen.Col Where the output pkeys are made available
936
    '''
937
    out_table = sql_gen.as_Table(out_table)
938
    mapping = sql_gen.ColDict(db, out_table, mapping)
939
    
940
    def log_debug(msg): db.log_debug(msg, level=1.5)
941
    def col_ustr(str_):
942
        return strings.repr_no_u(sql_gen.remove_col_rename(str_))
943
    
944
    log_debug('********** New iteration **********')
945
    log_debug('Inserting these input columns into '+strings.as_tt(
946
        out_table.to_str(db))+':\n'+strings.as_table(mapping, ustr=col_ustr))
947
    
948
    # Create input joins from list of input tables
949
    in_tables_ = in_tables[:] # don't modify input!
950
    in_tables0 = in_tables_.pop(0) # first table is separate
951
    in_pkey = pkey(db, in_tables0, recover=True)
952
    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
953
    input_joins = [in_tables0]+[sql_gen.Join(v,
954
        {in_pkey: sql_gen.join_same_not_null}) for v in in_tables_]
955
    
956
    if into == None:
957
        into = into_table_name(out_table, in_tables0, mapping, is_func)
958
    into = sql_gen.as_Table(into)
959
    
960
    log_debug('Joining together input tables into temp table')
961
    # Place in new table for speed and so don't modify input if values edited
962
    in_table = sql_gen.Table('in')
963
    flatten_cols = filter(sql_gen.is_table_col, mapping.values())
964
    mapping = dicts.join(mapping, flatten(db, in_table, input_joins,
965
        flatten_cols, preserve=[in_pkey_col], start=0))
966
    input_joins = [in_table]
967
    db.log_debug('Temp table: '+strings.as_tt(in_table.to_str(db)), level=2)
968
    
969
    # Resolve default value column
970
    try: default = mapping[default]
971
    except KeyError:
972
        if default != None:
973
            db.log_debug('Default value column '
974
                +strings.as_tt(strings.repr_no_u(default))
975
                +' does not exist in mapping, falling back to None', level=2.1)
976
            default = None
977
    
978
    out_pkey = pkey(db, out_table, recover=True)
979
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
980
    
981
    pkeys_names = [in_pkey, out_pkey]
982
    pkeys_cols = [in_pkey_col, out_pkey_col]
983
    
984
    pkeys_table_exists_ref = [False]
985
    def insert_into_pkeys(joins, cols):
986
        query, params = mk_select(db, joins, cols, order_by=None, start=0)
987
        if pkeys_table_exists_ref[0]:
988
            insert_select(db, into, pkeys_names, query, params)
989
        else:
990
            run_query_into(db, query, params, into=into)
991
            pkeys_table_exists_ref[0] = True
992
    
993
    limit_ref = [None]
994
    conds = set()
995
    distinct_on = []
996
    def mk_main_select(joins, cols):
997
        return mk_select(db, joins, cols, conds, distinct_on,
998
            limit=limit_ref[0], start=0)
999
    
1000
    exc_strs = set()
1001
    def log_exc(e):
1002
        e_str = exc.str_(e, first_line_only=True)
1003
        log_debug('Caught exception: '+e_str)
1004
        assert e_str not in exc_strs # avoid infinite loops
1005
        exc_strs.add(e_str)
1006
    def remove_all_rows():
1007
        log_debug('Returning NULL for all rows')
1008
        limit_ref[0] = 0 # just create an empty pkeys table
1009
    def ignore(in_col, value):
1010
        in_col_str = strings.as_tt(repr(in_col))
1011
        db.log_debug('Adding index on '+in_col_str+' to enable fast filtering',
1012
            level=2.5)
1013
        add_index(db, in_col)
1014
        log_debug('Ignoring rows with '+in_col_str+' = '
1015
            +strings.as_tt(repr(value)))
1016
    def remove_rows(in_col, value):
1017
        ignore(in_col, value)
1018
        cond = (in_col, sql_gen.CompareCond(value, '!='))
1019
        assert cond not in conds # avoid infinite loops
1020
        conds.add(cond)
1021
    def invalid2null(in_col, value):
1022
        ignore(in_col, value)
1023
        update(db, in_table, [(in_col, None)],
1024
            sql_gen.ColValueCond(in_col, value))
1025
    
1026
    def insert_pkeys_table(which):
1027
        return sql_gen.Table(sql_gen.add_suffix(in_table.name,
1028
            '_insert_'+which+'_pkeys'))
1029
    insert_out_pkeys = insert_pkeys_table('out')
1030
    insert_in_pkeys = insert_pkeys_table('in')
1031
    
1032
    # Do inserts and selects
1033
    join_cols = sql_gen.ColDict(db, out_table)
1034
    while True:
1035
        if limit_ref[0] == 0: # special case
1036
            log_debug('Creating an empty pkeys table')
1037
            cur = run_query_into(db, *mk_select(db, out_table, [out_pkey],
1038
                limit=limit_ref[0]), into=insert_out_pkeys)
1039
            break # don't do main case
1040
        
1041
        has_joins = join_cols != {}
1042
        
1043
        # Prepare to insert new rows
1044
        insert_joins = input_joins[:] # don't modify original!
1045
        insert_args = dict(recover=True, cacheable=False)
1046
        if has_joins:
1047
            distinct_on = [v.to_Col() for v in join_cols.values()]
1048
            insert_joins.append(sql_gen.Join(out_table, join_cols,
1049
                sql_gen.filter_out))
1050
        else:
1051
            insert_args.update(dict(returning=out_pkey, into=insert_out_pkeys))
1052
        main_select = mk_main_select(insert_joins, mapping.values())[0]
1053
        
1054
        log_debug('Trying to insert new rows')
1055
        try:
1056
            cur = insert_select(db, out_table, mapping.keys(), main_select,
1057
                **insert_args)
1058
            break # insert successful
1059
        except DuplicateKeyException, e:
1060
            log_exc(e)
1061
            
1062
            old_join_cols = join_cols.copy()
1063
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
1064
            log_debug('Ignoring existing rows, comparing on these columns:\n'
1065
                +strings.as_inline_table(join_cols, ustr=col_ustr))
1066
            assert join_cols != old_join_cols # avoid infinite loops
1067
        except NullValueException, e:
1068
            log_exc(e)
1069
            
1070
            out_col, = e.cols
1071
            try: in_col = mapping[out_col]
1072
            except KeyError:
1073
                log_debug('Missing mapping for NOT NULL column '+out_col)
1074
                remove_all_rows()
1075
            else: remove_rows(in_col, None)
1076
        except FunctionValueException, e:
1077
            log_exc(e)
1078
            
1079
            func_name = e.name
1080
            value = e.value
1081
            for out_col, in_col in mapping.iteritems():
1082
                invalid2null(sql_gen.unwrap_func_call(in_col, func_name), value)
1083
        except MissingCastException, e:
1084
            log_exc(e)
1085
            
1086
            out_col = e.col
1087
            mapping[out_col] = sql_gen.wrap_in_func(e.type, mapping[out_col])
1088
        except DatabaseErrors, e:
1089
            log_exc(e)
1090
            
1091
            msg = 'No handler for exception: '+exc.str_(e)
1092
            warnings.warn(DbWarning(msg))
1093
            log_debug(msg)
1094
            remove_all_rows()
1095
        # after exception handled, rerun loop with additional constraints
1096
    
1097
    if row_ct_ref != None and cur.rowcount >= 0:
1098
        row_ct_ref[0] += cur.rowcount
1099
    
1100
    if has_joins:
1101
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
1102
        log_debug('Getting output table pkeys of existing/inserted rows')
1103
        insert_into_pkeys(select_joins, pkeys_cols)
1104
    else:
1105
        add_row_num(db, insert_out_pkeys) # for joining with input pkeys
1106
        
1107
        log_debug('Getting input table pkeys of inserted rows')
1108
        run_query_into(db, *mk_main_select(input_joins, [in_pkey]),
1109
            into=insert_in_pkeys)
1110
        add_row_num(db, insert_in_pkeys) # for joining with output pkeys
1111
        
1112
        assert table_row_count(db, insert_out_pkeys) == table_row_count(db,
1113
            insert_in_pkeys)
1114
        
1115
        log_debug('Combining output and input pkeys in inserted order')
1116
        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,
1117
            {row_num_col: sql_gen.join_same_not_null})]
1118
        insert_into_pkeys(pkey_joins, pkeys_names)
1119
    
1120
    db.log_debug('Adding pkey on pkeys table to enable fast joins', level=2.5)
1121
    add_pkey(db, into)
1122
    
1123
    log_debug('Setting pkeys of missing rows to '+strings.as_tt(repr(default)))
1124
    missing_rows_joins = input_joins+[sql_gen.Join(into,
1125
        {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
1126
        # must use join_same_not_null or query will take forever
1127
    insert_into_pkeys(missing_rows_joins,
1128
        [in_pkey_col, sql_gen.NamedCol(out_pkey, default)])
1129
    
1130
    assert table_row_count(db, into) == table_row_count(db, in_table)
1131
    
1132
    return sql_gen.Col(out_pkey, into)
1133

    
1134
##### Data cleanup
1135

    
1136
def cleanup_table(db, table, cols):
1137
    def esc_name_(name): return esc_name(db, name)
1138
    
1139
    table = sql_gen.as_Table(table).to_str(db)
1140
    cols = map(esc_name_, cols)
1141
    
1142
    run_query(db, 'UPDATE '+table+' SET\n'+(',\n'.join(('\n'+col
1143
        +' = nullif(nullif(trim(both from '+col+"), %(null0)s), %(null1)s)"
1144
            for col in cols))),
1145
        dict(null0='', null1=r'\N'))
(24-24/36)