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):
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)
27

    
28
def _add_cursor_info(e, *args, **kw_args):
29
    '''For params, see get_cur_query()'''
30
    exc.add_msg(e, 'query: '+strings.ustr(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 ExceptionWithNameType(DbException):
50
    def __init__(self, type_, name, cause=None):
51
        DbException.__init__(self, 'for type: '+strings.as_tt(str(type_))
52
            +'; name: '+strings.as_tt(name), cause)
53
        self.type = type_
54
        self.name = name
55

    
56
class ConstraintException(DbException):
57
    def __init__(self, name, cols, cause=None):
58
        DbException.__init__(self, 'Violated '+strings.as_tt(name)
59
            +' constraint on columns: '+strings.as_tt(', '.join(cols)), cause)
60
        self.name = name
61
        self.cols = cols
62

    
63
class MissingCastException(DbException):
64
    def __init__(self, type_, col, cause=None):
65
        DbException.__init__(self, 'Missing cast to type '+strings.as_tt(type_)
66
            +' on column: '+strings.as_tt(col), cause)
67
        self.type = type_
68
        self.col = col
69

    
70
class NameException(DbException): pass
71

    
72
class DuplicateKeyException(ConstraintException): pass
73

    
74
class NullValueException(ConstraintException): pass
75

    
76
class FunctionValueException(ExceptionWithNameValue): pass
77

    
78
class DuplicateException(ExceptionWithNameType): pass
79

    
80
class EmptyRowException(DbException): pass
81

    
82
##### Warnings
83

    
84
class DbWarning(UserWarning): pass
85

    
86
##### Result retrieval
87

    
88
def col_names(cur): return (col[0] for col in cur.description)
89

    
90
def rows(cur): return iter(lambda: cur.fetchone(), None)
91

    
92
def consume_rows(cur):
93
    '''Used to fetch all rows so result will be cached'''
94
    iters.consume_iter(rows(cur))
95

    
96
def next_row(cur): return rows(cur).next()
97

    
98
def row(cur):
99
    row_ = next_row(cur)
100
    consume_rows(cur)
101
    return row_
102

    
103
def next_value(cur): return next_row(cur)[0]
104

    
105
def value(cur): return row(cur)[0]
106

    
107
def values(cur): return iters.func_iter(lambda: next_value(cur))
108

    
109
def value_or_none(cur):
110
    try: return value(cur)
111
    except StopIteration: return None
112

    
113
##### Escaping
114

    
115
def esc_name_by_module(module, name):
116
    if module == 'psycopg2' or module == None: quote = '"'
117
    elif module == 'MySQLdb': quote = '`'
118
    else: raise NotImplementedError("Can't escape name for "+module+' database')
119
    return sql_gen.esc_name(name, quote)
120

    
121
def esc_name_by_engine(engine, name, **kw_args):
122
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
123

    
124
def esc_name(db, name, **kw_args):
125
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
126

    
127
def qual_name(db, schema, table):
128
    def esc_name_(name): return esc_name(db, name)
129
    table = esc_name_(table)
130
    if schema != None: return esc_name_(schema)+'.'+table
131
    else: return table
132

    
133
##### Database connections
134

    
135
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
136

    
137
db_engines = {
138
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
139
    'PostgreSQL': ('psycopg2', {}),
140
}
141

    
142
DatabaseErrors_set = set([DbException])
143
DatabaseErrors = tuple(DatabaseErrors_set)
144

    
145
def _add_module(module):
146
    DatabaseErrors_set.add(module.DatabaseError)
147
    global DatabaseErrors
148
    DatabaseErrors = tuple(DatabaseErrors_set)
149

    
150
def db_config_str(db_config):
151
    return db_config['engine']+' database '+db_config['database']
152

    
153
log_debug_none = lambda msg, level=2: None
154

    
155
class DbConn:
156
    def __init__(self, db_config, autocommit=True, caching=True,
157
        log_debug=log_debug_none, debug_temp=False):
158
        '''
159
        @param debug_temp Whether temporary objects should instead be permanent.
160
            This assists in debugging the internal objects used by the program.
161
        '''
162
        self.db_config = db_config
163
        self.autocommit = autocommit
164
        self.caching = caching
165
        self.log_debug = log_debug
166
        self.debug = log_debug != log_debug_none
167
        self.debug_temp = debug_temp
168
        
169
        self.__db = None
170
        self.query_results = {}
171
        self._savepoint = 0
172
        self._notices_seen = set()
173
    
174
    def __getattr__(self, name):
175
        if name == '__dict__': raise Exception('getting __dict__')
176
        if name == 'db': return self._db()
177
        else: raise AttributeError()
178
    
179
    def __getstate__(self):
180
        state = copy.copy(self.__dict__) # shallow copy
181
        state['log_debug'] = None # don't pickle the debug callback
182
        state['_DbConn__db'] = None # don't pickle the connection
183
        return state
184
    
185
    def connected(self): return self.__db != None
186
    
187
    def _db(self):
188
        if self.__db == None:
189
            # Process db_config
190
            db_config = self.db_config.copy() # don't modify input!
191
            schemas = db_config.pop('schemas', None)
192
            module_name, mappings = db_engines[db_config.pop('engine')]
193
            module = __import__(module_name)
194
            _add_module(module)
195
            for orig, new in mappings.iteritems():
196
                try: util.rename_key(db_config, orig, new)
197
                except KeyError: pass
198
            
199
            # Connect
200
            self.__db = module.connect(**db_config)
201
            
202
            # Configure connection
203
            if hasattr(self.db, 'set_isolation_level'):
204
                import psycopg2.extensions
205
                self.db.set_isolation_level(
206
                    psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
207
            if schemas != None:
208
                search_path = [self.esc_name(s) for s in schemas.split(',')]
209
                search_path.append(value(run_query(self, 'SHOW search_path',
210
                    log_level=4)))
211
                run_query(self, 'SET search_path TO '+(','.join(search_path)),
212
                    log_level=3)
213
        
214
        return self.__db
215
    
216
    class DbCursor(Proxy):
217
        def __init__(self, outer):
218
            Proxy.__init__(self, outer.db.cursor())
219
            self.outer = outer
220
            self.query_results = outer.query_results
221
            self.query_lookup = None
222
            self.result = []
223
        
224
        def execute(self, query):
225
            self._is_insert = query.startswith('INSERT')
226
            self.query_lookup = query
227
            try:
228
                try:
229
                    cur = self.inner.execute(query)
230
                    self.outer.do_autocommit()
231
                finally: self.query = get_cur_query(self.inner, query)
232
            except Exception, e:
233
                _add_cursor_info(e, self, query)
234
                self.result = e # cache the exception as the result
235
                self._cache_result()
236
                raise
237
            
238
            # Always cache certain queries
239
            if query.startswith('CREATE') or query.startswith('ALTER'):
240
                # structural changes
241
                if query.find('ADD COLUMN') < 0: # rest of query must be unique
242
                    self._cache_result()
243
            elif self.rowcount == 0 and query.startswith('SELECT'): # empty
244
                consume_rows(self) # fetch all rows so result will be cached
245
            
246
            return cur
247
        
248
        def fetchone(self):
249
            row = self.inner.fetchone()
250
            if row != None: self.result.append(row)
251
            # otherwise, fetched all rows
252
            else: self._cache_result()
253
            return row
254
        
255
        def _cache_result(self):
256
            # For inserts that return a result set, don't cache result set since
257
            # inserts are not idempotent. Other non-SELECT queries don't have
258
            # their result set read, so only exceptions will be cached (an
259
            # invalid query will always be invalid).
260
            if self.query_results != None and (not self._is_insert
261
                or isinstance(self.result, Exception)):
262
                
263
                assert self.query_lookup != None
264
                self.query_results[self.query_lookup] = self.CacheCursor(
265
                    util.dict_subset(dicts.AttrsDictView(self),
266
                    ['query', 'result', 'rowcount', 'description']))
267
        
268
        class CacheCursor:
269
            def __init__(self, cached_result): self.__dict__ = cached_result
270
            
271
            def execute(self, *args, **kw_args):
272
                if isinstance(self.result, Exception): raise self.result
273
                # otherwise, result is a rows list
274
                self.iter = iter(self.result)
275
            
276
            def fetchone(self):
277
                try: return self.iter.next()
278
                except StopIteration: return None
279
    
280
    def esc_value(self, value):
281
        try: str_ = self.mogrify('%s', [value])
282
        except NotImplementedError, e:
283
            module = util.root_module(self.db)
284
            if module == 'MySQLdb':
285
                import _mysql
286
                str_ = _mysql.escape_string(value)
287
            else: raise e
288
        return strings.to_unicode(str_)
289
    
290
    def esc_name(self, name): return esc_name(self, name) # calls global func
291
    
292
    def std_code(self, str_):
293
        '''Standardizes SQL code.
294
        * Ensures that string literals are prefixed by `E`
295
        '''
296
        if str_.startswith("'"): str_ = 'E'+str_
297
        return str_
298
    
299
    def can_mogrify(self):
300
        module = util.root_module(self.db)
301
        return module == 'psycopg2'
302
    
303
    def mogrify(self, query, params=None):
304
        if self.can_mogrify(): return self.db.cursor().mogrify(query, params)
305
        else: raise NotImplementedError("Can't mogrify query")
306
    
307
    def print_notices(self):
308
        if hasattr(self.db, 'notices'):
309
            for msg in self.db.notices:
310
                if msg not in self._notices_seen:
311
                    self._notices_seen.add(msg)
312
                    self.log_debug(msg, level=2)
313
    
314
    def run_query(self, query, cacheable=False, log_level=2,
315
        debug_msg_ref=None):
316
        '''
317
        @param log_ignore_excs The log_level will be increased by 2 if the query
318
            throws one of these exceptions.
319
        @param debug_msg_ref If specified, the log message will be returned in
320
            this instead of being output. This allows you to filter log messages
321
            depending on the result of the query.
322
        '''
323
        assert query != None
324
        
325
        if not self.caching: cacheable = False
326
        used_cache = False
327
        
328
        def log_msg(query):
329
            if used_cache: cache_status = 'cache hit'
330
            elif cacheable: cache_status = 'cache miss'
331
            else: cache_status = 'non-cacheable'
332
            return 'DB query: '+cache_status+':\n'+strings.as_code(query, 'SQL')
333
        
334
        try:
335
            # Get cursor
336
            if cacheable:
337
                try:
338
                    cur = self.query_results[query]
339
                    used_cache = True
340
                except KeyError: cur = self.DbCursor(self)
341
            else: cur = self.db.cursor()
342
            
343
            # Log query
344
            if self.debug and debug_msg_ref == None: # log before running
345
                self.log_debug(log_msg(query), log_level)
346
            
347
            # Run query
348
            cur.execute(query)
349
        finally:
350
            self.print_notices()
351
            if self.debug and debug_msg_ref != None: # return after running
352
                debug_msg_ref[0] = log_msg(str(get_cur_query(cur, query)))
353
        
354
        return cur
355
    
356
    def is_cached(self, query): return query in self.query_results
357
    
358
    def with_autocommit(self, func):
359
        import psycopg2.extensions
360
        
361
        prev_isolation_level = self.db.isolation_level
362
        self.db.set_isolation_level(
363
            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
364
        try: return func()
365
        finally: self.db.set_isolation_level(prev_isolation_level)
366
    
367
    def with_savepoint(self, func):
368
        savepoint = 'level_'+str(self._savepoint)
369
        self.run_query('SAVEPOINT '+savepoint, log_level=4)
370
        self._savepoint += 1
371
        try: return func()
372
        except:
373
            self.run_query('ROLLBACK TO SAVEPOINT '+savepoint, log_level=4)
374
            raise
375
        finally:
376
            # Always release savepoint, because after ROLLBACK TO SAVEPOINT,
377
            # "The savepoint remains valid and can be rolled back to again"
378
            # (http://www.postgresql.org/docs/8.3/static/sql-rollback-to.html).
379
            self.run_query('RELEASE SAVEPOINT '+savepoint, log_level=4)
380
            
381
            self._savepoint -= 1
382
            assert self._savepoint >= 0
383
            
384
            self.do_autocommit() # OK to do this after ROLLBACK TO SAVEPOINT
385
    
386
    def do_autocommit(self):
387
        '''Autocommits if outside savepoint'''
388
        assert self._savepoint >= 0
389
        if self.autocommit and self._savepoint == 0:
390
            self.log_debug('Autocommitting', level=4)
391
            self.db.commit()
392
    
393
    def col_info(self, col):
394
        table = sql_gen.Table('columns', 'information_schema')
395
        cols = ['data_type', 'column_default',
396
            cast(self, 'boolean', 'is_nullable')]
397
        
398
        conds = [('table_name', col.table.name), ('column_name', col.name)]
399
        schema = col.table.schema
400
        if schema != None: conds.append(('table_schema', schema))
401
        
402
        type_, default, nullable = row(select(self, table, cols, conds,
403
            order_by='table_schema', limit=1, log_level=4))
404
            # TODO: order_by search_path schema order
405
        default = sql_gen.as_Code(default, self)
406
        
407
        return sql_gen.TypedCol(col.name, type_, default, nullable)
408
    
409
    def TempFunction(self, name):
410
        if self.debug_temp: schema = None
411
        else: schema = 'pg_temp'
412
        return sql_gen.Function(name, schema)
413

    
414
connect = DbConn
415

    
416
##### Recoverable querying
417

    
418
def with_savepoint(db, func): return db.with_savepoint(func)
419

    
420
def run_query(db, query, recover=None, cacheable=False, log_level=2,
421
    log_ignore_excs=None, **kw_args):
422
    '''For params, see DbConn.run_query()'''
423
    if recover == None: recover = False
424
    if log_ignore_excs == None: log_ignore_excs = ()
425
    log_ignore_excs = tuple(log_ignore_excs)
426
    
427
    debug_msg_ref = None # usually, db.run_query() logs query before running it
428
    # But if filtering with log_ignore_excs, wait until after exception parsing
429
    if log_ignore_excs != () or not db.can_mogrify(): debug_msg_ref = [None]
430
    
431
    try:
432
        try:
433
            def run(): return db.run_query(query, cacheable, log_level,
434
                debug_msg_ref, **kw_args)
435
            if recover and not db.is_cached(query):
436
                return with_savepoint(db, run)
437
            else: return run() # don't need savepoint if cached
438
        except Exception, e:
439
            msg = exc.str_(e)
440
            
441
            match = re.search(r'duplicate key value violates unique constraint '
442
                r'"((_?[^\W_]+)_.+?)"', msg)
443
            if match:
444
                constraint, table = match.groups()
445
                cols = []
446
                if recover: # need auto-rollback to run index_cols()
447
                    try: cols = index_cols(db, table, constraint)
448
                    except NotImplementedError: pass
449
                raise DuplicateKeyException(constraint, cols, e)
450
            
451
            match = re.search(r'null value in column "(.+?)" violates not-null'
452
                r' constraint', msg)
453
            if match: raise NullValueException('NOT NULL', [match.group(1)], e)
454
            
455
            match = re.search(r'\b(?:invalid input (?:syntax|value)\b.*?'
456
                r'|date/time field value out of range): "(.+?)"\n'
457
                r'(?:(?s).*?)\bfunction "(.+?)"', msg)
458
            if match:
459
                value, name = match.groups()
460
                raise FunctionValueException(name, strings.to_unicode(value), e)
461
            
462
            match = re.search(r'column "(.+?)" is of type (.+?) but expression '
463
                r'is of type', msg)
464
            if match:
465
                col, type_ = match.groups()
466
                raise MissingCastException(type_, col, e)
467
            
468
            match = re.search(r'\b(\S+) "(.+?)".*? already exists', msg)
469
            if match:
470
                type_, name = match.groups()
471
                raise DuplicateException(type_, name, e)
472
            
473
            raise # no specific exception raised
474
    except log_ignore_excs:
475
        log_level += 2
476
        raise
477
    finally:
478
        if debug_msg_ref != None and debug_msg_ref[0] != None:
479
            db.log_debug(debug_msg_ref[0], log_level)
480

    
481
##### Basic queries
482

    
483
def next_version(name):
484
    version = 1 # first existing name was version 0
485
    match = re.match(r'^(.*)#(\d+)$', name)
486
    if match:
487
        name, version = match.groups()
488
        version = int(version)+1
489
    return sql_gen.concat(name, '#'+str(version))
490

    
491
def lock_table(db, table, mode):
492
    table = sql_gen.as_Table(table)
493
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
494

    
495
def run_query_into(db, query, into=None, add_indexes_=False, **kw_args):
496
    '''Outputs a query to a temp table.
497
    For params, see run_query().
498
    '''
499
    if into == None: return run_query(db, query, **kw_args)
500
    
501
    assert isinstance(into, sql_gen.Table)
502
    
503
    into.is_temp = True
504
    # "temporary tables cannot specify a schema name", so remove schema
505
    into.schema = None
506
    
507
    kw_args['recover'] = True
508
    kw_args.setdefault('log_ignore_excs', (DuplicateException,))
509
    
510
    temp = not db.debug_temp # tables are permanent in debug_temp mode
511
    
512
    # Create table
513
    while True:
514
        create_query = 'CREATE'
515
        if temp: create_query += ' TEMP'
516
        create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
517
        
518
        try:
519
            cur = run_query(db, create_query, **kw_args)
520
                # CREATE TABLE AS sets rowcount to # rows in query
521
            break
522
        except DuplicateException, e:
523
            into.name = next_version(into.name)
524
            # try again with next version of name
525
    
526
    if add_indexes_: add_indexes(db, into)
527
    
528
    return cur
529

    
530
order_by_pkey = object() # tells mk_select() to order by the pkey
531

    
532
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
533

    
534
def mk_select(db, tables, fields=None, conds=None, distinct_on=[], limit=None,
535
    start=None, order_by=order_by_pkey, default_table=None):
536
    '''
537
    @param tables The single table to select from, or a list of tables to join
538
        together, with tables after the first being sql_gen.Join objects
539
    @param fields Use None to select all fields in the table
540
    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
541
        * container can be any iterable type
542
        * compare_left_side: sql_gen.Code|str (for col name)
543
        * compare_right_side: sql_gen.ValueCond|literal value
544
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
545
        use all columns
546
    @return query
547
    '''
548
    # Parse tables param
549
    tables = lists.mk_seq(tables)
550
    tables = list(tables) # don't modify input! (list() copies input)
551
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
552
    
553
    # Parse other params
554
    if conds == None: conds = []
555
    elif dicts.is_dict(conds): conds = conds.items()
556
    conds = list(conds) # don't modify input! (list() copies input)
557
    assert limit == None or type(limit) == int
558
    assert start == None or type(start) == int
559
    if order_by is order_by_pkey:
560
        if distinct_on != []: order_by = None
561
        else: order_by = pkey(db, table0, recover=True)
562
    
563
    query = 'SELECT'
564
    
565
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
566
    
567
    # DISTINCT ON columns
568
    if distinct_on != []:
569
        query += '\nDISTINCT'
570
        if distinct_on is not distinct_on_all:
571
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
572
    
573
    # Columns
574
    if fields == None:
575
        if query.find('\n') >= 0: whitespace = '\n'
576
        else: whitespace = ' '
577
        query += whitespace+'*'
578
    else:
579
        assert fields != []
580
        query += '\n'+('\n, '.join(map(parse_col, fields)))
581
    
582
    # Main table
583
    query += '\nFROM '+table0.to_str(db)
584
    
585
    # Add joins
586
    left_table = table0
587
    for join_ in tables:
588
        table = join_.table
589
        
590
        # Parse special values
591
        if join_.type_ is sql_gen.filter_out: # filter no match
592
            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
593
                sql_gen.CompareCond(None, '~=')))
594
        
595
        query += '\n'+join_.to_str(db, left_table)
596
        
597
        left_table = table
598
    
599
    missing = True
600
    if conds != []:
601
        if len(conds) == 1: whitespace = ' '
602
        else: whitespace = '\n'
603
        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
604
            .to_str(db) for l, r in conds], 'WHERE')
605
        missing = False
606
    if order_by != None:
607
        query += '\nORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)
608
    if limit != None: query += '\nLIMIT '+str(limit); missing = False
609
    if start != None:
610
        if start != 0: query += '\nOFFSET '+str(start)
611
        missing = False
612
    if missing: warnings.warn(DbWarning(
613
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
614
    
615
    return query
616

    
617
def select(db, *args, **kw_args):
618
    '''For params, see mk_select() and run_query()'''
619
    recover = kw_args.pop('recover', None)
620
    cacheable = kw_args.pop('cacheable', True)
621
    log_level = kw_args.pop('log_level', 2)
622
    
623
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
624
        log_level=log_level)
625

    
626
def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
627
    embeddable=False, ignore=False):
628
    '''
629
    @param returning str|None An inserted column (such as pkey) to return
630
    @param embeddable Whether the query should be embeddable as a nested SELECT.
631
        Warning: If you set this and cacheable=True when the query is run, the
632
        query will be fully cached, not just if it raises an exception.
633
    @param ignore Whether to ignore duplicate keys.
634
    '''
635
    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
636
    if cols == []: cols = None # no cols (all defaults) = unknown col names
637
    if cols != None: cols = [sql_gen.to_name_only_col(c, table) for c in cols]
638
    if select_query == None: select_query = 'DEFAULT VALUES'
639
    if returning != None: returning = sql_gen.as_Col(returning, table)
640
    
641
    first_line = 'INSERT INTO '+table.to_str(db)
642
    
643
    def mk_insert(select_query):
644
        query = first_line
645
        if cols != None:
646
            query += '\n('+(', '.join((c.to_str(db) for c in cols)))+')'
647
        query += '\n'+select_query
648
        
649
        if returning != None:
650
            returning_name_col = sql_gen.to_name_only_col(returning)
651
            query += '\nRETURNING '+returning_name_col.to_str(db)
652
        
653
        return query
654
    
655
    return_type = 'unknown'
656
    if returning != None: return_type = returning.to_str(db)+'%TYPE'
657
    
658
    lang = 'sql'
659
    if ignore:
660
        assert cols != None
661
        # Always return something to set the correct rowcount
662
        if returning == None: returning = sql_gen.NamedCol('NULL', None)
663
        
664
        embeddable = True # must use function
665
        lang = 'plpgsql'
666
        row = [sql_gen.Col(c.name, 'row') for c in cols]
667
        
668
        query = '''\
669
DECLARE
670
    row '''+table.to_str(db)+'''%ROWTYPE;
671
    cur CURSOR FOR
672
'''+select_query+'''
673
;
674
BEGIN
675
    OPEN cur;
676
    
677
    /* Need an EXCEPTION block for each individual row because "When an error is
678
    caught by an EXCEPTION clause, [...] all changes to persistent database
679
    state within the block are rolled back."
680
    This is unfortunate because "A block containing an EXCEPTION clause is
681
    significantly more expensive to enter and exit than a block without one."
682
    (http://www.postgresql.org/docs/8.3/static/plpgsql-control-structures.html\
683
#PLPGSQL-ERROR-TRAPPING)
684
    */
685
    WHILE true LOOP
686
        BEGIN
687
            FETCH FROM cur INTO \
688
'''+(', '.join((c.to_str(db) for c in row)))+''';
689
            EXIT WHEN NOT FOUND;
690
            
691
            RETURN QUERY
692
'''+mk_insert(sql_gen.Values(row).to_str(db))+'''
693
;
694
        EXCEPTION
695
            /* Note that cursor stays at current position in spite of automatic
696
            ROLLBACK to before EXCEPTION block: "If a previously opened cursor
697
            is affected by a FETCH command inside a savepoint that is later
698
            rolled back, the cursor position remains at the position that FETCH
699
            left it pointing to (that is, FETCH is not rolled back)."
700
            (http://www.postgresql.org/docs/8.3/static/sql-rollback-to.html)
701
            */
702
            WHEN unique_violation THEN NULL; -- continue to next row
703
        END;
704
    END LOOP;
705
    
706
    CLOSE cur;
707
END;\
708
'''
709
    else: query = mk_insert(select_query)
710
    
711
    if embeddable:
712
        # Create function
713
        function_name = sql_gen.clean_name(first_line)
714
        while True:
715
            try:
716
                function = db.TempFunction(function_name)
717
                
718
                function_query = '''\
719
CREATE FUNCTION '''+function.to_str(db)+'''()
720
RETURNS SETOF '''+return_type+'''
721
LANGUAGE '''+lang+'''
722
AS $$
723
'''+query+'''
724
$$;
725
'''
726
                run_query(db, function_query, recover=True, cacheable=True,
727
                    log_ignore_excs=(DuplicateException,))
728
                break # this version was successful
729
            except DuplicateException, e:
730
                function_name = next_version(function_name)
731
                # try again with next version of name
732
        
733
        # Return query that uses function
734
        cols = None
735
        if returning != None: cols = [returning]
736
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(function),
737
            cols) # AS clause requires function alias
738
        return mk_select(db, func_table, start=0, order_by=None)
739
    
740
    return query
741

    
742
def insert_select(db, *args, **kw_args):
743
    '''For params, see mk_insert_select() and run_query_into()
744
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
745
        values in
746
    '''
747
    into = kw_args.pop('into', None)
748
    if into != None: kw_args['embeddable'] = True
749
    recover = kw_args.pop('recover', None)
750
    if kw_args.get('ignore', False): recover = True
751
    cacheable = kw_args.pop('cacheable', True)
752
    log_level = kw_args.pop('log_level', 2)
753
    
754
    return run_query_into(db, mk_insert_select(db, *args, **kw_args), into,
755
        recover=recover, cacheable=cacheable, log_level=log_level)
756

    
757
default = sql_gen.default # tells insert() to use the default value for a column
758

    
759
def insert(db, table, row, *args, **kw_args):
760
    '''For params, see insert_select()'''
761
    if lists.is_seq(row): cols = None
762
    else:
763
        cols = row.keys()
764
        row = row.values()
765
    row = list(row) # ensure that "== []" works
766
    
767
    if row == []: query = None
768
    else: query = sql_gen.Values(row).to_str(db)
769
    
770
    return insert_select(db, table, cols, query, *args, **kw_args)
771

    
772
def mk_update(db, table, changes=None, cond=None):
773
    '''
774
    @param changes [(col, new_value),...]
775
        * container can be any iterable type
776
        * col: sql_gen.Code|str (for col name)
777
        * new_value: sql_gen.Code|literal value
778
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
779
    @return str query
780
    '''
781
    query = 'UPDATE '+sql_gen.as_Table(table).to_str(db)+'\nSET\n'
782
    query += ',\n'.join((sql_gen.to_name_only_col(col, table).to_str(db)+' = '
783
        +sql_gen.as_Value(new_value).to_str(db) for col, new_value in changes))
784
    if cond != None: query += '\nWHERE\n'+cond.to_str(db)
785
    
786
    return query
787

    
788
def update(db, *args, **kw_args):
789
    '''For params, see mk_update() and run_query()'''
790
    recover = kw_args.pop('recover', None)
791
    log_level = kw_args.pop('log_level', 2)
792
    
793
    return run_query(db, mk_update(db, *args, **kw_args), recover,
794
        log_level=log_level)
795

    
796
def last_insert_id(db):
797
    module = util.root_module(db.db)
798
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
799
    elif module == 'MySQLdb': return db.insert_id()
800
    else: return None
801

    
802
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
803
    '''Creates a mapping from original column names (which may have collisions)
804
    to names that will be distinct among the columns' tables.
805
    This is meant to be used for several tables that are being joined together.
806
    @param cols The columns to combine. Duplicates will be removed.
807
    @param into The table for the new columns.
808
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
809
        columns will be included in the mapping even if they are not in cols.
810
        The tables of the provided Col objects will be changed to into, so make
811
        copies of them if you want to keep the original tables.
812
    @param as_items Whether to return a list of dict items instead of a dict
813
    @return dict(orig_col=new_col, ...)
814
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
815
        * new_col: sql_gen.Col(orig_col_name, into)
816
        * All mappings use the into table so its name can easily be
817
          changed for all columns at once
818
    '''
819
    cols = lists.uniqify(cols)
820
    
821
    items = []
822
    for col in preserve:
823
        orig_col = copy.copy(col)
824
        col.table = into
825
        items.append((orig_col, col))
826
    preserve = set(preserve)
827
    for col in cols:
828
        if col not in preserve:
829
            items.append((col, sql_gen.Col(str(col), into, col.srcs)))
830
    
831
    if not as_items: items = dict(items)
832
    return items
833

    
834
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
835
    '''For params, see mk_flatten_mapping()
836
    @return See return value of mk_flatten_mapping()
837
    '''
838
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
839
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
840
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
841
        into=into, add_indexes_=True)
842
    return dict(items)
843

    
844
def mk_track_data_error(db, errors_table, cols, value, error_code, error):
845
    assert cols != ()
846
    
847
    cols = map(sql_gen.to_name_only_col, cols)
848
    
849
    columns_cols = ['column']
850
    columns = sql_gen.NamedValues('columns', columns_cols,
851
        [[c.name] for c in cols])
852
    values_cols = ['value', 'error_code', 'error']
853
    values = sql_gen.NamedValues('values', values_cols,
854
        [value, error_code, error])
855
    
856
    select_cols = columns_cols+values_cols
857
    name_only_cols = map(sql_gen.to_name_only_col, select_cols)
858
    errors_table = sql_gen.NamedTable('errors', errors_table)
859
    joins = [columns, sql_gen.Join(values, type_='CROSS'),
860
        sql_gen.Join(errors_table, dict(zip(name_only_cols, select_cols)),
861
        sql_gen.filter_out)]
862
    
863
    return mk_insert_select(db, errors_table, name_only_cols,
864
        mk_select(db, joins, select_cols, order_by=None))
865

    
866
def track_data_error(db, errors_table, cols, *args, **kw_args):
867
    '''
868
    @param errors_table If None, does nothing.
869
    '''
870
    if errors_table == None or cols == (): return
871
    run_query(db, mk_track_data_error(db, errors_table, cols, *args, **kw_args),
872
        cacheable=True, log_level=4)
873

    
874
def cast(db, type_, col, errors_table=None):
875
    '''Casts an (unrenamed) column or value.
876
    If errors_table set and col has srcs, saves errors in errors_table (using
877
    col's srcs attr as the source columns) and converts errors to warnings.
878
    @param col str|sql_gen.Col|sql_gen.Literal
879
    @param errors_table None|sql_gen.Table|str
880
    '''
881
    col = sql_gen.as_Col(col)
882
    save_errors = (errors_table != None and isinstance(col, sql_gen.Col)
883
        and col.srcs != ())
884
    if not save_errors: return sql_gen.Cast(type_, col) # can't save errors
885
    
886
    assert not isinstance(col, sql_gen.NamedCol)
887
    
888
    errors_table = sql_gen.as_Table(errors_table)
889
    srcs = map(sql_gen.to_name_only_col, col.srcs)
890
    function_name = str(sql_gen.FunctionCall(type_, *srcs))
891
    function = db.TempFunction(function_name)
892
    
893
    while True:
894
        # Create function definition
895
        query = '''\
896
CREATE FUNCTION '''+function.to_str(db)+'''(value text)
897
RETURNS '''+type_+'''
898
LANGUAGE plpgsql
899
STRICT
900
AS $$
901
BEGIN
902
    /* The explicit cast to the return type is needed to make the cast happen
903
    inside the try block. (Implicit casts to the return type happen at the end
904
    of the function, outside any block.) */
905
    RETURN value::'''+type_+''';
906
EXCEPTION
907
    WHEN data_exception THEN
908
        -- Save error in errors table.
909
        -- Insert the value and error for *each* source column.
910
'''+mk_track_data_error(db, errors_table, srcs,
911
    *map(sql_gen.CustomCode, ['value', 'SQLSTATE', 'SQLERRM']))+''';
912
        
913
        RAISE WARNING '%', SQLERRM;
914
        RETURN NULL;
915
END;
916
$$;
917
'''
918
        
919
        # Create function
920
        try:
921
            run_query(db, query, recover=True, cacheable=True,
922
                log_ignore_excs=(DuplicateException,))
923
            break # successful
924
        except DuplicateException:
925
            function.name = next_version(function.name)
926
            # try again with next version of name
927
    
928
    return sql_gen.FunctionCall(function, col)
929

    
930
##### Database structure queries
931

    
932
def table_row_count(db, table, recover=None):
933
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
934
        order_by=None, start=0), recover=recover, log_level=3))
935

    
936
def table_cols(db, table, recover=None):
937
    return list(col_names(select(db, table, limit=0, order_by=None,
938
        recover=recover, log_level=4)))
939

    
940
def pkey(db, table, recover=None):
941
    '''Assumed to be first column in table'''
942
    return table_cols(db, table, recover)[0]
943

    
944
not_null_col = 'not_null_col'
945

    
946
def table_not_null_col(db, table, recover=None):
947
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
948
    if not_null_col in table_cols(db, table, recover): return not_null_col
949
    else: return pkey(db, table, recover)
950

    
951
def index_cols(db, table, index):
952
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
953
    automatically created. When you don't know whether something is a UNIQUE
954
    constraint or a UNIQUE index, use this function.'''
955
    module = util.root_module(db.db)
956
    if module == 'psycopg2':
957
        return list(values(run_query(db, '''\
958
SELECT attname
959
FROM
960
(
961
        SELECT attnum, attname
962
        FROM pg_index
963
        JOIN pg_class index ON index.oid = indexrelid
964
        JOIN pg_class table_ ON table_.oid = indrelid
965
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
966
        WHERE
967
            table_.relname = '''+db.esc_value(table)+'''
968
            AND index.relname = '''+db.esc_value(index)+'''
969
    UNION
970
        SELECT attnum, attname
971
        FROM
972
        (
973
            SELECT
974
                indrelid
975
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
976
                    AS indkey
977
            FROM pg_index
978
            JOIN pg_class index ON index.oid = indexrelid
979
            JOIN pg_class table_ ON table_.oid = indrelid
980
            WHERE
981
                table_.relname = '''+db.esc_value(table)+'''
982
                AND index.relname = '''+db.esc_value(index)+'''
983
        ) s
984
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
985
) s
986
ORDER BY attnum
987
'''
988
            , cacheable=True, log_level=4)))
989
    else: raise NotImplementedError("Can't list index columns for "+module+
990
        ' database')
991

    
992
def constraint_cols(db, table, constraint):
993
    module = util.root_module(db.db)
994
    if module == 'psycopg2':
995
        return list(values(run_query(db, '''\
996
SELECT attname
997
FROM pg_constraint
998
JOIN pg_class ON pg_class.oid = conrelid
999
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
1000
WHERE
1001
    relname = '''+db.esc_value(table)+'''
1002
    AND conname = '''+db.esc_value(constraint)+'''
1003
ORDER BY attnum
1004
'''
1005
            )))
1006
    else: raise NotImplementedError("Can't list constraint columns for "+module+
1007
        ' database')
1008

    
1009
row_num_col = '_row_num'
1010

    
1011
def add_index(db, exprs, table=None, unique=False, ensure_not_null_=True):
1012
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
1013
    Currently, only function calls are supported as expressions.
1014
    @param ensure_not_null_ If set, translates NULL values to sentinel values.
1015
        This allows indexes to be used for comparisons where NULLs are equal.
1016
    '''
1017
    exprs = lists.mk_seq(exprs)
1018
    
1019
    # Parse exprs
1020
    old_exprs = exprs[:]
1021
    exprs = []
1022
    cols = []
1023
    for i, expr in enumerate(old_exprs):
1024
        expr = sql_gen.as_Col(expr, table)
1025
        
1026
        # Handle nullable columns
1027
        if ensure_not_null_:
1028
            try: expr = ensure_not_null(db, expr)
1029
            except KeyError: pass # unknown type, so just create plain index
1030
        
1031
        # Extract col
1032
        expr = copy.deepcopy(expr) # don't modify input!
1033
        if isinstance(expr, sql_gen.FunctionCall):
1034
            col = expr.args[0]
1035
            expr = sql_gen.Expr(expr)
1036
        else: col = expr
1037
        assert isinstance(col, sql_gen.Col)
1038
        
1039
        # Extract table
1040
        if table == None:
1041
            assert sql_gen.is_table_col(col)
1042
            table = col.table
1043
        
1044
        col.table = None
1045
        
1046
        exprs.append(expr)
1047
        cols.append(col)
1048
    
1049
    table = sql_gen.as_Table(table)
1050
    index = sql_gen.Table(str(sql_gen.Col(','.join(map(str, cols)), table)))
1051
    
1052
    # Add index
1053
    while True:
1054
        str_ = 'CREATE'
1055
        if unique: str_ += ' UNIQUE'
1056
        str_ += ' INDEX '+index.to_str(db)+' ON '+table.to_str(db)+' ('+(
1057
            ', '.join((v.to_str(db) for v in exprs)))+')'
1058
        
1059
        try:
1060
            run_query(db, str_, recover=True, cacheable=True, log_level=3,
1061
                log_ignore_excs=(DuplicateException,))
1062
            break
1063
        except DuplicateException:
1064
            index.name = next_version(index.name)
1065
            # try again with next version of name
1066

    
1067
def add_pkey(db, table, cols=None, recover=None):
1068
    '''Adds a primary key.
1069
    @param cols [sql_gen.Col,...] The columns in the primary key.
1070
        Defaults to the first column in the table.
1071
    @pre The table must not already have a primary key.
1072
    '''
1073
    table = sql_gen.as_Table(table)
1074
    if cols == None: cols = [pkey(db, table, recover)]
1075
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
1076
    
1077
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
1078
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
1079
        log_ignore_excs=(DuplicateException,))
1080

    
1081
def add_not_null(db, col):
1082
    table = col.table
1083
    col = sql_gen.to_name_only_col(col)
1084
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1085
        +col.to_str(db)+' SET NOT NULL', cacheable=True)
1086

    
1087
def add_index_col(db, col, suffix, expr, nullable=True):
1088
    if sql_gen.index_col(col) != None: return # already has index col
1089
    
1090
    new_col = sql_gen.suffixed_col(col, suffix)
1091
    
1092
    # Add column
1093
    while True:
1094
        new_typed_col = sql_gen.TypedCol(new_col.name, db.col_info(col).type)
1095
        try:
1096
            add_col(db, col.table, new_typed_col,
1097
                log_ignore_excs=(DuplicateException,))
1098
            break
1099
        except DuplicateException:
1100
            new_col.name = next_version(new_col.name)
1101
            # try again with next version of name
1102
    
1103
    update(db, col.table, [(new_col, expr)])
1104
    if not nullable: add_not_null(db, new_col)
1105
    add_index(db, new_col)
1106
    
1107
    col.table.index_cols[col.name] = new_col
1108

    
1109
def ensure_not_null(db, col):
1110
    '''For params, see sql_gen.ensure_not_null()'''
1111
    expr = sql_gen.ensure_not_null(db, col)
1112
    
1113
    # If nullable column in a temp table, add separate column instead
1114
    if sql_gen.is_temp_col(col) and isinstance(expr, sql_gen.EnsureNotNull):
1115
        add_index_col(db, col, '::NOT NULL', expr, nullable=False)
1116
        expr = sql_gen.index_col(col)
1117
    
1118
    return expr
1119

    
1120
already_indexed = object() # tells add_indexes() the pkey has already been added
1121

    
1122
def add_indexes(db, table, has_pkey=True):
1123
    '''Adds an index on all columns in a table.
1124
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
1125
        index should be added on the first column.
1126
        * If already_indexed, the pkey is assumed to have already been added
1127
    '''
1128
    cols = table_cols(db, table)
1129
    if has_pkey:
1130
        if has_pkey is not already_indexed: add_pkey(db, table)
1131
        cols = cols[1:]
1132
    for col in cols: add_index(db, col, table)
1133

    
1134
def add_col(db, table, col, **kw_args):
1135
    assert isinstance(col, sql_gen.TypedCol)
1136
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '
1137
        +col.to_str(db), recover=True, cacheable=True, **kw_args)
1138

    
1139
row_num_typed_col = sql_gen.TypedCol(row_num_col, 'serial', nullable=False,
1140
    constraints='PRIMARY KEY')
1141

    
1142
def add_row_num(db, table):
1143
    '''Adds a row number column to a table. Its name is in row_num_col. It will
1144
    be the primary key.'''
1145
    add_col(db, table, row_num_typed_col, log_level=3)
1146

    
1147
def cast_temp_col(db, type_, col, errors_table=None):
1148
    '''Like cast(), but creates a new column with the cast values if the input
1149
    is a column.
1150
    @return The new column or cast value
1151
    '''
1152
    def cast_(col): return cast(db, type_, col, errors_table)
1153
    
1154
    try: col = sql_gen.underlying_col(col)
1155
    except sql_gen.NoUnderlyingTableException: return sql_gen.wrap(cast_, col)
1156
    
1157
    table = col.table
1158
    new_col = sql_gen.Col(sql_gen.concat(col.name, '::'+type_), table, col.srcs)
1159
    expr = cast_(col)
1160
    add_col(db, table, sql_gen.TypedCol(new_col.name, type_))
1161
    update(db, table, [(new_col, expr)])
1162
    add_index(db, new_col)
1163
    
1164
    return new_col
1165

    
1166
def drop_table(db, table):
1167
    table = sql_gen.as_Table(table)
1168
    return run_query(db, 'DROP TABLE IF EXISTS '+table.to_str(db)+' CASCADE')
1169

    
1170
def create_table(db, table, cols, has_pkey=True, col_indexes=True):
1171
    '''Creates a table.
1172
    @param cols [sql_gen.TypedCol,...] The column names and types
1173
    @param has_pkey If set, the first column becomes the primary key.
1174
    @param col_indexes bool|[ref]
1175
        * If True, indexes will be added on all non-pkey columns.
1176
        * If a list reference, [0] will be set to a function to do this.
1177
          This can be used to delay index creation until the table is populated.
1178
    '''
1179
    table = sql_gen.as_Table(table)
1180
    
1181
    if has_pkey:
1182
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1183
        pkey.constraints = 'PRIMARY KEY'
1184
    
1185
    str_ = 'CREATE TABLE '+table.to_str(db)+' (\n'
1186
    str_ += '\n, '.join(v.to_str(db) for v in cols)
1187
    str_ += '\n);\n'
1188
    run_query(db, str_, cacheable=True, log_level=2)
1189
    
1190
    # Add indexes
1191
    if has_pkey: has_pkey = already_indexed
1192
    def add_indexes_(): add_indexes(db, table, has_pkey)
1193
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1194
    elif col_indexes: add_indexes_() # add now
1195

    
1196
def vacuum(db, table):
1197
    table = sql_gen.as_Table(table)
1198
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1199
        log_level=3))
1200

    
1201
def truncate(db, table, schema='public', **kw_args):
1202
    '''For params, see run_query()'''
1203
    table = sql_gen.as_Table(table, schema)
1204
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE', **kw_args)
1205

    
1206
def empty_temp(db, tables):
1207
    if db.debug_temp: return # leave temp tables there for debugging
1208
    tables = lists.mk_seq(tables)
1209
    for table in tables: truncate(db, table, log_level=3)
1210

    
1211
def tables(db, schema_like='public', table_like='%', exact=False):
1212
    if exact: compare = '='
1213
    else: compare = 'LIKE'
1214
    
1215
    module = util.root_module(db.db)
1216
    if module == 'psycopg2':
1217
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1218
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1219
        return values(select(db, 'pg_tables', ['tablename'], conds,
1220
            order_by='tablename', log_level=4))
1221
    elif module == 'MySQLdb':
1222
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1223
            , cacheable=True, log_level=4))
1224
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1225

    
1226
def table_exists(db, table):
1227
    table = sql_gen.as_Table(table)
1228
    return list(tables(db, table.schema, table.name, exact=True)) != []
1229

    
1230
def function_exists(db, function):
1231
    function = sql_gen.as_Function(function)
1232
    
1233
    info_table = sql_gen.Table('routines', 'information_schema')
1234
    conds = [('routine_name', function.name)]
1235
    schema = function.schema
1236
    if schema != None: conds.append(('routine_schema', schema))
1237
    # Exclude trigger functions, since they cannot be called directly
1238
    conds.append(('data_type', sql_gen.CompareCond('trigger', '!=')))
1239
    
1240
    return list(values(select(db, info_table, ['routine_name'], conds,
1241
        order_by='routine_schema', limit=1, log_level=4))) != []
1242
        # TODO: order_by search_path schema order
1243

    
1244
def errors_table(db, table, if_exists=True):
1245
    '''
1246
    @param if_exists If set, returns None if the errors table doesn't exist
1247
    @return None|sql_gen.Table
1248
    '''
1249
    table = sql_gen.as_Table(table)
1250
    if table.srcs != (): table = table.srcs[0]
1251
    
1252
    errors_table = sql_gen.suffixed_table(table, '.errors')
1253
    if if_exists and not table_exists(db, errors_table): return None
1254
    return errors_table
1255

    
1256
##### Database management
1257

    
1258
def empty_db(db, schema='public', **kw_args):
1259
    '''For kw_args, see tables()'''
1260
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1261

    
1262
##### Heuristic queries
1263

    
1264
def put(db, table, row, pkey_=None, row_ct_ref=None):
1265
    '''Recovers from errors.
1266
    Only works under PostgreSQL (uses INSERT RETURNING).
1267
    '''
1268
    row = sql_gen.ColDict(db, table, row)
1269
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
1270
    
1271
    try:
1272
        cur = insert(db, table, row, pkey_, recover=True)
1273
        if row_ct_ref != None and cur.rowcount >= 0:
1274
            row_ct_ref[0] += cur.rowcount
1275
        return value(cur)
1276
    except DuplicateKeyException, e:
1277
        row = sql_gen.ColDict(db, table,
1278
            util.dict_subset_right_join(row, e.cols))
1279
        return value(select(db, table, [pkey_], row, recover=True))
1280

    
1281
def get(db, table, row, pkey, row_ct_ref=None, create=False):
1282
    '''Recovers from errors'''
1283
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
1284
    except StopIteration:
1285
        if not create: raise
1286
        return put(db, table, row, pkey, row_ct_ref) # insert new row
1287

    
1288
def is_func_result(col):
1289
    return col.table.name.find('(') >= 0 and col.name == 'result'
1290

    
1291
def into_table_name(out_table, in_tables0, mapping, is_func):
1292
    def in_col_str(in_col):
1293
        in_col = sql_gen.remove_col_rename(in_col)
1294
        if isinstance(in_col, sql_gen.Col):
1295
            table = in_col.table
1296
            if table == in_tables0:
1297
                in_col = sql_gen.to_name_only_col(in_col)
1298
            elif is_func_result(in_col): in_col = table # omit col name
1299
        return str(in_col)
1300
    
1301
    str_ = str(out_table)
1302
    if is_func:
1303
        str_ += '('
1304
        
1305
        try: value_in_col = mapping['value']
1306
        except KeyError:
1307
            str_ += ', '.join((str(k)+'='+in_col_str(v)
1308
                for k, v in mapping.iteritems()))
1309
        else: str_ += in_col_str(value_in_col)
1310
        
1311
        str_ += ')'
1312
    else:
1313
        out_col = 'rank'
1314
        try: in_col = mapping[out_col]
1315
        except KeyError: str_ += '_pkeys'
1316
        else: # has a rank column, so hierarchical
1317
            str_ += '['+str(out_col)+'='+in_col_str(in_col)+']'
1318
    return str_
1319

    
1320
def put_table(db, out_table, in_tables, mapping, row_ct_ref=None, into=None,
1321
    default=None, is_func=False, on_error=exc.raise_):
1322
    '''Recovers from errors.
1323
    Only works under PostgreSQL (uses INSERT RETURNING).
1324
    IMPORTANT: Must be run at the *beginning* of a transaction.
1325
    @param in_tables The main input table to select from, followed by a list of
1326
        tables to join with it using the main input table's pkey
1327
    @param mapping dict(out_table_col=in_table_col, ...)
1328
        * out_table_col: str (*not* sql_gen.Col)
1329
        * in_table_col: sql_gen.Col|literal-value
1330
    @param into The table to contain the output and input pkeys.
1331
        Defaults to `out_table.name+'_pkeys'`.
1332
    @param default The *output* column to use as the pkey for missing rows.
1333
        If this output column does not exist in the mapping, uses None.
1334
    @param is_func Whether out_table is the name of a SQL function, not a table
1335
    @return sql_gen.Col Where the output pkeys are made available
1336
    '''
1337
    out_table = sql_gen.as_Table(out_table)
1338
    
1339
    def log_debug(msg): db.log_debug(msg, level=1.5)
1340
    def col_ustr(str_):
1341
        return strings.repr_no_u(sql_gen.remove_col_rename(str_))
1342
    
1343
    log_debug('********** New iteration **********')
1344
    log_debug('Inserting these input columns into '+strings.as_tt(
1345
        out_table.to_str(db))+':\n'+strings.as_table(mapping, ustr=col_ustr))
1346
    
1347
    is_function = function_exists(db, out_table)
1348
    
1349
    if is_function: out_pkey = 'result'
1350
    else: out_pkey = pkey(db, out_table, recover=True)
1351
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
1352
    
1353
    if mapping == {}: # need at least one column for INSERT SELECT
1354
        mapping = {out_pkey: None} # ColDict will replace with default value
1355
    
1356
    # Create input joins from list of input tables
1357
    in_tables_ = in_tables[:] # don't modify input!
1358
    in_tables0 = in_tables_.pop(0) # first table is separate
1359
    errors_table_ = errors_table(db, in_tables0)
1360
    in_pkey = pkey(db, in_tables0, recover=True)
1361
    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
1362
    input_joins = [in_tables0]+[sql_gen.Join(v,
1363
        {in_pkey: sql_gen.join_same_not_null}) for v in in_tables_]
1364
    
1365
    if into == None:
1366
        into = into_table_name(out_table, in_tables0, mapping, is_func)
1367
    into = sql_gen.as_Table(into)
1368
    
1369
    # Set column sources
1370
    in_cols = filter(sql_gen.is_table_col, mapping.values())
1371
    for col in in_cols:
1372
        if col.table == in_tables0: col.set_srcs(sql_gen.src_self)
1373
    
1374
    log_debug('Joining together input tables into temp table')
1375
    # Place in new table for speed and so don't modify input if values edited
1376
    in_table = sql_gen.Table('in')
1377
    mapping = dicts.join(mapping, flatten(db, in_table, input_joins, in_cols,
1378
        preserve=[in_pkey_col], start=0))
1379
    input_joins = [in_table]
1380
    db.log_debug('Temp table: '+strings.as_tt(in_table.to_str(db)), level=2)
1381
    
1382
    mapping = sql_gen.ColDict(db, out_table, mapping)
1383
        # after applying dicts.join() because that returns a plain dict
1384
    
1385
    # Resolve default value column
1386
    if default != None:
1387
        try: default = mapping[default]
1388
        except KeyError:
1389
            db.log_debug('Default value column '
1390
                +strings.as_tt(strings.repr_no_u(default))
1391
                +' does not exist in mapping, falling back to None', level=2.1)
1392
            default = None
1393
    
1394
    pkeys_names = [in_pkey, out_pkey]
1395
    pkeys_cols = [in_pkey_col, out_pkey_col]
1396
    
1397
    pkeys_table_exists_ref = [False]
1398
    def insert_into_pkeys(joins, cols, distinct=False):
1399
        kw_args = {}
1400
        if distinct: kw_args.update(dict(distinct_on=[in_pkey_col]))
1401
        query = mk_select(db, joins, cols, order_by=None, start=0, **kw_args)
1402
        
1403
        if pkeys_table_exists_ref[0]:
1404
            insert_select(db, into, pkeys_names, query)
1405
        else:
1406
            run_query_into(db, query, into=into)
1407
            pkeys_table_exists_ref[0] = True
1408
    
1409
    limit_ref = [None]
1410
    conds = set()
1411
    distinct_on = sql_gen.ColDict(db, out_table)
1412
    def mk_main_select(joins, cols):
1413
        distinct_on_cols = [c.to_Col() for c in distinct_on.values()]
1414
        return mk_select(db, joins, cols, conds, distinct_on_cols,
1415
            limit=limit_ref[0], start=0)
1416
    
1417
    exc_strs = set()
1418
    def log_exc(e):
1419
        e_str = exc.str_(e, first_line_only=True)
1420
        log_debug('Caught exception: '+e_str)
1421
        assert e_str not in exc_strs # avoid infinite loops
1422
        exc_strs.add(e_str)
1423
    
1424
    def remove_all_rows():
1425
        log_debug('Ignoring all rows')
1426
        limit_ref[0] = 0 # just create an empty pkeys table
1427
    
1428
    def ignore(in_col, value, e):
1429
        track_data_error(db, errors_table_, in_col.srcs, value, e.cause.pgcode,
1430
            e.cause.pgerror)
1431
        log_debug('Ignoring rows with '+strings.as_tt(repr(in_col))+' = '
1432
            +strings.as_tt(repr(value)))
1433
    
1434
    def remove_rows(in_col, value, e):
1435
        ignore(in_col, value, e)
1436
        cond = (in_col, sql_gen.CompareCond(value, '!='))
1437
        assert cond not in conds # avoid infinite loops
1438
        conds.add(cond)
1439
    
1440
    def invalid2null(in_col, value, e):
1441
        ignore(in_col, value, e)
1442
        update(db, in_table, [(in_col, None)],
1443
            sql_gen.ColValueCond(in_col, value))
1444
    
1445
    def insert_pkeys_table(which):
1446
        return sql_gen.Table(sql_gen.concat(in_table.name,
1447
            '_insert_'+which+'_pkeys'))
1448
    insert_out_pkeys = insert_pkeys_table('out')
1449
    insert_in_pkeys = insert_pkeys_table('in')
1450
    
1451
    # Do inserts and selects
1452
    join_cols = sql_gen.ColDict(db, out_table)
1453
    while True:
1454
        if limit_ref[0] == 0: # special case
1455
            log_debug('Creating an empty pkeys table')
1456
            cur = run_query_into(db, mk_select(db, out_table, [out_pkey],
1457
                limit=limit_ref[0]), into=insert_out_pkeys)
1458
            break # don't do main case
1459
        
1460
        has_joins = join_cols != {}
1461
        
1462
        log_debug('Trying to insert new rows')
1463
        
1464
        # Prepare to insert new rows
1465
        insert_joins = input_joins[:] # don't modify original!
1466
        insert_args = dict(recover=True, cacheable=False)
1467
        if has_joins:
1468
            insert_args.update(dict(ignore=True))
1469
        else:
1470
            insert_args.update(dict(returning=out_pkey, into=insert_out_pkeys))
1471
        main_select = mk_main_select(insert_joins, mapping.values())
1472
        
1473
        def main_insert():
1474
            if is_function:
1475
                log_debug('Calling function on input rows')
1476
                args = dict(((k.name, v) for k, v in mapping.iteritems()))
1477
                func_call = sql_gen.NamedCol(out_pkey,
1478
                    sql_gen.FunctionCall(out_table, **args))
1479
                insert_into_pkeys(input_joins, [in_pkey_col, func_call])
1480
                return None
1481
            else:
1482
                lock_table(db, out_table, 'EXCLUSIVE')
1483
                return insert_select(db, out_table, mapping.keys(), main_select,
1484
                    **insert_args)
1485
        
1486
        try:
1487
            cur = with_savepoint(db, main_insert)
1488
            break # insert successful
1489
        except MissingCastException, e:
1490
            log_exc(e)
1491
            
1492
            out_col = e.col
1493
            type_ = e.type
1494
            
1495
            log_debug('Casting '+strings.as_tt(out_col)+' input to '
1496
                +strings.as_tt(type_))
1497
            mapping[out_col] = cast_temp_col(db, type_, mapping[out_col],
1498
                errors_table_)
1499
        except DuplicateKeyException, e:
1500
            log_exc(e)
1501
            
1502
            old_join_cols = join_cols.copy()
1503
            distinct_on.update(util.dict_subset(mapping, e.cols))
1504
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
1505
            log_debug('Ignoring existing rows, comparing on these columns:\n'
1506
                +strings.as_inline_table(join_cols, ustr=col_ustr))
1507
            assert join_cols != old_join_cols # avoid infinite loops
1508
        except NullValueException, e:
1509
            log_exc(e)
1510
            
1511
            out_col, = e.cols
1512
            try: in_col = mapping[out_col]
1513
            except KeyError:
1514
                log_debug('Missing mapping for NOT NULL column '+out_col)
1515
                remove_all_rows()
1516
            else: remove_rows(in_col, None, e)
1517
        except FunctionValueException, e:
1518
            log_exc(e)
1519
            
1520
            func_name = e.name
1521
            value = e.value
1522
            for out_col, in_col in mapping.iteritems():
1523
                in_col = sql_gen.unwrap_func_call(in_col, func_name)
1524
                invalid2null(in_col, value, e)
1525
        except DatabaseErrors, e:
1526
            log_exc(e)
1527
            
1528
            log_debug('No handler for exception')
1529
            on_error(e)
1530
            remove_all_rows()
1531
        # after exception handled, rerun loop with additional constraints
1532
    
1533
    if cur != None and row_ct_ref != None and cur.rowcount >= 0:
1534
        row_ct_ref[0] += cur.rowcount
1535
    
1536
    if is_function: pass # pkeys table already created
1537
    elif has_joins:
1538
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
1539
        log_debug('Getting output table pkeys of existing/inserted rows')
1540
        insert_into_pkeys(select_joins, pkeys_cols, distinct=True)
1541
    else:
1542
        add_row_num(db, insert_out_pkeys) # for joining with input pkeys
1543
        
1544
        log_debug('Getting input table pkeys of inserted rows')
1545
        run_query_into(db, mk_main_select(input_joins, [in_pkey]),
1546
            into=insert_in_pkeys)
1547
        add_row_num(db, insert_in_pkeys) # for joining with output pkeys
1548
        
1549
        assert table_row_count(db, insert_out_pkeys) == table_row_count(db,
1550
            insert_in_pkeys)
1551
        
1552
        log_debug('Combining output and input pkeys in inserted order')
1553
        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,
1554
            {row_num_col: sql_gen.join_same_not_null})]
1555
        insert_into_pkeys(pkey_joins, pkeys_names)
1556
        
1557
        empty_temp(db, [insert_out_pkeys, insert_in_pkeys])
1558
    
1559
    db.log_debug('Adding pkey on pkeys table to enable fast joins', level=2.5)
1560
    add_pkey(db, into)
1561
    
1562
    log_debug('Setting pkeys of missing rows to '+strings.as_tt(repr(default)))
1563
    missing_rows_joins = input_joins+[sql_gen.Join(into,
1564
        {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
1565
        # must use join_same_not_null or query will take forever
1566
    insert_into_pkeys(missing_rows_joins,
1567
        [in_pkey_col, sql_gen.NamedCol(out_pkey, default)])
1568
    
1569
    assert table_row_count(db, into) == table_row_count(db, in_table)
1570
    
1571
    empty_temp(db, in_table)
1572
    
1573
    srcs = []
1574
    if is_func: srcs = sql_gen.cols_srcs(in_cols)
1575
    return sql_gen.Col(out_pkey, into, srcs)
1576

    
1577
##### Data cleanup
1578

    
1579
def cleanup_table(db, table, cols):
1580
    table = sql_gen.as_Table(table)
1581
    cols = map(sql_gen.as_Col, cols)
1582
    
1583
    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '
1584
        +db.esc_value(r'\N')+')')
1585
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))
1586
        for v in cols]
1587
    
1588
    update(db, table, changes)
(24-24/36)