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)
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
connect = DbConn
333

    
334
##### Querying
335

    
336
def run_raw_query(db, *args, **kw_args):
337
    '''For params, see DbConn.run_query()'''
338
    return db.run_query(*args, **kw_args)
339

    
340
def mogrify(db, query, params):
341
    module = util.root_module(db.db)
342
    if module == 'psycopg2': return db.db.cursor().mogrify(query, params)
343
    else: raise NotImplementedError("Can't mogrify query for "+module+
344
        ' database')
345

    
346
##### Recoverable querying
347

    
348
def with_savepoint(db, func): return db.with_savepoint(func)
349

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

    
407
##### Basic queries
408

    
409
def next_version(name):
410
    version = 1 # first existing name was version 0
411
    match = re.match(r'^(.*)#(\d+)$', name)
412
    if match:
413
        name, version = match.groups()
414
        version = int(version)+1
415
    return sql_gen.add_suffix(name, '#'+str(version))
416

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

    
444
order_by_pkey = object() # tells mk_select() to order by the pkey
445

    
446
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
447

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

    
527
def select(db, *args, **kw_args):
528
    '''For params, see mk_select() and run_query()'''
529
    recover = kw_args.pop('recover', None)
530
    cacheable = kw_args.pop('cacheable', True)
531
    log_level = kw_args.pop('log_level', 2)
532
    
533
    query, params = mk_select(db, *args, **kw_args)
534
    return run_query(db, query, params, recover, cacheable, log_level=log_level)
535

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

    
596
def insert_select(db, *args, **kw_args):
597
    '''For params, see mk_insert_select() and run_query_into()
598
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
599
        values in
600
    '''
601
    into = kw_args.pop('into', None)
602
    if into != None: kw_args['embeddable'] = True
603
    recover = kw_args.pop('recover', None)
604
    cacheable = kw_args.pop('cacheable', True)
605
    
606
    query, params = mk_insert_select(db, *args, **kw_args)
607
    return run_query_into(db, query, params, into, recover=recover,
608
        cacheable=cacheable)
609

    
610
default = object() # tells insert() to use the default value for a column
611

    
612
def insert(db, table, row, *args, **kw_args):
613
    '''For params, see insert_select()'''
614
    if lists.is_seq(row): cols = None
615
    else:
616
        cols = row.keys()
617
        row = row.values()
618
    row = list(row) # ensure that "!= []" works
619
    
620
    # Check for special values
621
    labels = []
622
    values = []
623
    for value in row:
624
        if value is default: labels.append('DEFAULT')
625
        else:
626
            labels.append('%s')
627
            values.append(value)
628
    
629
    # Build query
630
    if values != []: query = 'VALUES ('+(', '.join(labels))+')'
631
    else: query = None
632
    
633
    return insert_select(db, table, cols, query, values, *args, **kw_args)
634

    
635
def mk_update(db, table, changes=None, cond=None):
636
    '''
637
    @param changes [(col, new_value),...]
638
        * container can be any iterable type
639
        * col: sql_gen.Code|str (for col name)
640
        * new_value: sql_gen.Code|literal value
641
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
642
    @return str query
643
    '''
644
    query = 'UPDATE '+sql_gen.as_Table(table).to_str(db)+'\nSET\n'
645
    query += ',\n'.join((sql_gen.to_name_only_col(col, table).to_str(db)+' = '
646
        +sql_gen.as_Value(new_value).to_str(db) for col, new_value in changes))
647
    if cond != None: query += '\nWHERE\n'+cond.to_str(db)
648
    
649
    return query
650

    
651
def update(db, *args, **kw_args):
652
    '''For params, see mk_update() and run_query()'''
653
    recover = kw_args.pop('recover', None)
654
    
655
    return run_query(db, mk_update(db, *args, **kw_args), [], recover)
656

    
657
def last_insert_id(db):
658
    module = util.root_module(db.db)
659
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
660
    elif module == 'MySQLdb': return db.insert_id()
661
    else: return None
662

    
663
def truncate(db, table, schema='public'):
664
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
665

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

    
697
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
698
    '''For params, see mk_flatten_mapping()
699
    @return See return value of mk_flatten_mapping()
700
    '''
701
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
702
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
703
    run_query_into(db, *mk_select(db, joins, cols, limit=limit, start=start),
704
        into=into)
705
    return dict(items)
706

    
707
##### Database structure queries
708

    
709
def table_row_count(db, table, recover=None):
710
    return value(run_query(db, *mk_select(db, table, [sql_gen.row_count],
711
        order_by=None, start=0), recover=recover, log_level=3))
712

    
713
def table_cols(db, table, recover=None):
714
    return list(col_names(select(db, table, limit=0, order_by=None,
715
        recover=recover, log_level=4)))
716

    
717
def pkey(db, table, recover=None):
718
    '''Assumed to be first column in table'''
719
    return table_cols(db, table, recover)[0]
720

    
721
not_null_col = 'not_null_col'
722

    
723
def table_not_null_col(db, table, recover=None):
724
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
725
    if not_null_col in table_cols(db, table, recover): return not_null_col
726
    else: return pkey(db, table, recover)
727

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

    
769
def constraint_cols(db, table, constraint):
770
    module = util.root_module(db.db)
771
    if module == 'psycopg2':
772
        return list(values(run_query(db, '''\
773
SELECT attname
774
FROM pg_constraint
775
JOIN pg_class ON pg_class.oid = conrelid
776
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
777
WHERE
778
    relname = %(table)s
779
    AND conname = %(constraint)s
780
ORDER BY attnum
781
''',
782
            {'table': table, 'constraint': constraint})))
783
    else: raise NotImplementedError("Can't list constraint columns for "+module+
784
        ' database')
785

    
786
row_num_col = '_row_num'
787

    
788
def add_index(db, expr):
789
    '''Adds an index on a column or expression if it doesn't already exist.
790
    Currently, only function calls are supported as expressions.
791
    '''
792
    expr = copy.copy(expr) # don't modify input!
793
    
794
    # Extract col
795
    if isinstance(expr, sql_gen.FunctionCall):
796
        col = expr.args[0]
797
        expr = sql_gen.Expr(expr)
798
    else: col = expr
799
    assert sql_gen.is_table_col(col)
800
    
801
    index = sql_gen.as_Table(str(expr))
802
    table = col.table
803
    col.table = None
804
    try: run_query(db, 'CREATE INDEX '+index.to_str(db)+' ON '+table.to_str(db)
805
        +' ('+expr.to_str(db)+')', recover=True, cacheable=True, log_level=3)
806
    except DuplicateTableException: pass # index already existed
807

    
808
def add_pkey(db, table, recover=None):
809
    '''Makes the first column in a table the primary key.
810
    @pre The table must not already have a primary key.
811
    '''
812
    table = sql_gen.as_Table(table)
813
    
814
    index = sql_gen.as_Table(sql_gen.add_suffix(table.name, '_pkey'))
815
    col = sql_gen.to_name_only_col(pkey(db, table, recover))
816
    try:
817
        run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD CONSTRAINT '
818
            +index.to_str(db)+' PRIMARY KEY('+col.to_str(db)+')',
819
            recover=True, cacheable=True, log_level=3,
820
            log_ignore_excs=(DuplicateTableException,))
821
    except DuplicateTableException, e:
822
        index.name = next_version(index.name)
823
        # try again with next version of name
824

    
825
def add_row_num(db, table):
826
    '''Adds a row number column to a table. Its name is in row_num_col. It will
827
    be the primary key.'''
828
    table = sql_gen.as_Table(table).to_str(db)
829
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
830
        +' serial NOT NULL PRIMARY KEY', log_level=3)
831

    
832
def tables(db, schema_like='public', table_like='%'):
833
    module = util.root_module(db.db)
834
    params = {'schema_like': schema_like, 'table_like': table_like}
835
    if module == 'psycopg2':
836
        return values(run_query(db, '''\
837
SELECT tablename
838
FROM pg_tables
839
WHERE
840
    schemaname LIKE %(schema_like)s
841
    AND tablename LIKE %(table_like)s
842
ORDER BY tablename
843
''',
844
            params, cacheable=True))
845
    elif module == 'MySQLdb':
846
        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,
847
            cacheable=True))
848
    else: raise NotImplementedError("Can't list tables for "+module+' database')
849

    
850
##### Database management
851

    
852
def empty_db(db, schema='public', **kw_args):
853
    '''For kw_args, see tables()'''
854
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
855

    
856
##### Heuristic queries
857

    
858
def put(db, table, row, pkey_=None, row_ct_ref=None):
859
    '''Recovers from errors.
860
    Only works under PostgreSQL (uses INSERT RETURNING).
861
    '''
862
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
863
    
864
    try:
865
        cur = insert(db, table, row, pkey_, recover=True)
866
        if row_ct_ref != None and cur.rowcount >= 0:
867
            row_ct_ref[0] += cur.rowcount
868
        return value(cur)
869
    except DuplicateKeyException, e:
870
        return value(select(db, table, [pkey_],
871
            util.dict_subset_right_join(row, e.cols), recover=True))
872

    
873
def get(db, table, row, pkey, row_ct_ref=None, create=False):
874
    '''Recovers from errors'''
875
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
876
    except StopIteration:
877
        if not create: raise
878
        return put(db, table, row, pkey, row_ct_ref) # insert new row
879

    
880
def is_func_result(col):
881
    return col.table.name.find('(') >= 0 and col.name == 'result'
882

    
883
def into_table_name(out_table, in_tables0, mapping, is_func):
884
    def in_col_str(in_col):
885
        in_col = sql_gen.remove_col_rename(in_col)
886
        if isinstance(in_col, sql_gen.Col):
887
            table = in_col.table
888
            if table == in_tables0:
889
                in_col = sql_gen.to_name_only_col(in_col)
890
            elif is_func_result(in_col): in_col = table # omit col name
891
        return str(in_col)
892
    
893
    str_ = str(out_table)
894
    if is_func:
895
        str_ += '('
896
        
897
        try: value_in_col = mapping['value']
898
        except KeyError:
899
            str_ += ', '.join((k+'='+in_col_str(v)
900
                for k, v in mapping.iteritems()))
901
        else: str_ += in_col_str(value_in_col)
902
        
903
        str_ += ')'
904
    else: str_ += '_pkeys'
905
    return str_
906

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

    
1120
##### Data cleanup
1121

    
1122
def cleanup_table(db, table, cols):
1123
    def esc_name_(name): return esc_name(db, name)
1124
    
1125
    table = sql_gen.as_Table(table).to_str(db)
1126
    cols = map(esc_name_, cols)
1127
    
1128
    run_query(db, 'UPDATE '+table+' SET\n'+(',\n'.join(('\n'+col
1129
        +' = nullif(nullif(trim(both from '+col+"), %(null0)s), %(null1)s)"
1130
            for col in cols))),
1131
        dict(null0='', null1=r'\N'))
(24-24/36)