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
            if self.rowcount == 0 and query.startswith('SELECT'): # empty SELECT
238
                consume_rows(self) # fetch all rows so result will be cached
239
            return cur
240
        
241
        def fetchone(self):
242
            row = self.inner.fetchone()
243
            if row != None: self.result.append(row)
244
            # otherwise, fetched all rows
245
            else: self._cache_result()
246
            return row
247
        
248
        def _cache_result(self):
249
            # For inserts that return a result set, don't cache result set since
250
            # inserts are not idempotent. Other non-SELECT queries don't have
251
            # their result set read, so only exceptions will be cached (an
252
            # invalid query will always be invalid).
253
            if self.query_results != None and (not self._is_insert
254
                or isinstance(self.result, Exception)):
255
                
256
                assert self.query_lookup != None
257
                self.query_results[self.query_lookup] = self.CacheCursor(
258
                    util.dict_subset(dicts.AttrsDictView(self),
259
                    ['query', 'result', 'rowcount', 'description']))
260
        
261
        class CacheCursor:
262
            def __init__(self, cached_result): self.__dict__ = cached_result
263
            
264
            def execute(self, *args, **kw_args):
265
                if isinstance(self.result, Exception): raise self.result
266
                # otherwise, result is a rows list
267
                self.iter = iter(self.result)
268
            
269
            def fetchone(self):
270
                try: return self.iter.next()
271
                except StopIteration: return None
272
    
273
    def esc_value(self, value):
274
        try: str_ = self.mogrify('%s', [value])
275
        except NotImplementedError, e:
276
            module = util.root_module(self.db)
277
            if module == 'MySQLdb':
278
                import _mysql
279
                str_ = _mysql.escape_string(value)
280
            else: raise e
281
        return strings.to_unicode(str_)
282
    
283
    def esc_name(self, name): return esc_name(self, name) # calls global func
284
    
285
    def std_code(self, str_):
286
        '''Standardizes SQL code.
287
        * Ensures that string literals are prefixed by `E`
288
        '''
289
        if str_.startswith("'"): str_ = 'E'+str_
290
        return str_
291
    
292
    def can_mogrify(self):
293
        module = util.root_module(self.db)
294
        return module == 'psycopg2'
295
    
296
    def mogrify(self, query, params=None):
297
        if self.can_mogrify(): return self.db.cursor().mogrify(query, params)
298
        else: raise NotImplementedError("Can't mogrify query")
299
    
300
    def print_notices(self):
301
        if hasattr(self.db, 'notices'):
302
            for msg in self.db.notices:
303
                if msg not in self._notices_seen:
304
                    self._notices_seen.add(msg)
305
                    self.log_debug(msg, level=2)
306
    
307
    def run_query(self, query, cacheable=False, log_level=2,
308
        debug_msg_ref=None):
309
        '''
310
        @param log_ignore_excs The log_level will be increased by 2 if the query
311
            throws one of these exceptions.
312
        @param debug_msg_ref If specified, the log message will be returned in
313
            this instead of being output. This allows you to filter log messages
314
            depending on the result of the query.
315
        '''
316
        assert query != None
317
        
318
        if not self.caching: cacheable = False
319
        used_cache = False
320
        
321
        def log_msg(query):
322
            if used_cache: cache_status = 'cache hit'
323
            elif cacheable: cache_status = 'cache miss'
324
            else: cache_status = 'non-cacheable'
325
            return 'DB query: '+cache_status+':\n'+strings.as_code(query, 'SQL')
326
        
327
        try:
328
            # Get cursor
329
            if cacheable:
330
                try:
331
                    cur = self.query_results[query]
332
                    used_cache = True
333
                except KeyError: cur = self.DbCursor(self)
334
            else: cur = self.db.cursor()
335
            
336
            # Log query
337
            if self.debug and debug_msg_ref == None: # log before running
338
                self.log_debug(log_msg(query), log_level)
339
            
340
            # Run query
341
            cur.execute(query)
342
        finally:
343
            self.print_notices()
344
            if self.debug and debug_msg_ref != None: # return after running
345
                debug_msg_ref[0] = log_msg(str(get_cur_query(cur, query)))
346
        
347
        return cur
348
    
349
    def is_cached(self, query): return query in self.query_results
350
    
351
    def with_autocommit(self, func):
352
        import psycopg2.extensions
353
        
354
        prev_isolation_level = self.db.isolation_level
355
        self.db.set_isolation_level(
356
            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
357
        try: return func()
358
        finally: self.db.set_isolation_level(prev_isolation_level)
359
    
360
    def with_savepoint(self, func):
361
        savepoint = 'level_'+str(self._savepoint)
362
        self.run_query('SAVEPOINT '+savepoint, log_level=4)
363
        self._savepoint += 1
364
        try: return func()
365
        except:
366
            self.run_query('ROLLBACK TO SAVEPOINT '+savepoint, log_level=4)
367
            raise
368
        finally:
369
            # Always release savepoint, because after ROLLBACK TO SAVEPOINT,
370
            # "The savepoint remains valid and can be rolled back to again"
371
            # (http://www.postgresql.org/docs/8.3/static/sql-rollback-to.html).
372
            self.run_query('RELEASE SAVEPOINT '+savepoint, log_level=4)
373
            
374
            self._savepoint -= 1
375
            assert self._savepoint >= 0
376
            
377
            self.do_autocommit() # OK to do this after ROLLBACK TO SAVEPOINT
378
    
379
    def do_autocommit(self):
380
        '''Autocommits if outside savepoint'''
381
        assert self._savepoint >= 0
382
        if self.autocommit and self._savepoint == 0:
383
            self.log_debug('Autocommitting', level=4)
384
            self.db.commit()
385
    
386
    def col_info(self, col):
387
        table = sql_gen.Table('columns', 'information_schema')
388
        cols = ['data_type', 'column_default',
389
            cast(self, 'boolean', 'is_nullable')]
390
        
391
        conds = [('table_name', col.table.name), ('column_name', col.name)]
392
        schema = col.table.schema
393
        if schema != None: conds.append(('table_schema', schema))
394
        
395
        type_, default, nullable = row(select(self, table, cols, conds,
396
            order_by='table_schema', limit=1, log_level=4))
397
            # TODO: order_by search_path schema order
398
        default = sql_gen.as_Code(default, self)
399
        
400
        return sql_gen.TypedCol(col.name, type_, default, nullable)
401
    
402
    def TempFunction(self, name):
403
        if self.debug_temp: schema = None
404
        else: schema = 'pg_temp'
405
        return sql_gen.Function(name, schema)
406

    
407
connect = DbConn
408

    
409
##### Recoverable querying
410

    
411
def with_savepoint(db, func): return db.with_savepoint(func)
412

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

    
473
##### Basic queries
474

    
475
def next_version(name):
476
    version = 1 # first existing name was version 0
477
    match = re.match(r'^(.*)#(\d+)$', name)
478
    if match:
479
        name, version = match.groups()
480
        version = int(version)+1
481
    return sql_gen.concat(name, '#'+str(version))
482

    
483
def lock_table(db, table, mode):
484
    table = sql_gen.as_Table(table)
485
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
486

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

    
522
order_by_pkey = object() # tells mk_select() to order by the pkey
523

    
524
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
525

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

    
607
def select(db, *args, **kw_args):
608
    '''For params, see mk_select() and run_query()'''
609
    recover = kw_args.pop('recover', None)
610
    cacheable = kw_args.pop('cacheable', True)
611
    log_level = kw_args.pop('log_level', 2)
612
    
613
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
614
        log_level=log_level)
615

    
616
def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
617
    embeddable=False):
618
    '''
619
    @param returning str|None An inserted column (such as pkey) to return
620
    @param embeddable Whether the query should be embeddable as a nested SELECT.
621
        Warning: If you set this and cacheable=True when the query is run, the
622
        query will be fully cached, not just if it raises an exception.
623
    '''
624
    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
625
    if cols == []: cols = None # no cols (all defaults) = unknown col names
626
    if cols != None:
627
        cols = [sql_gen.to_name_only_col(v, table).to_str(db) for v in cols]
628
    if select_query == None: select_query = 'DEFAULT VALUES'
629
    if returning != None: returning = sql_gen.as_Col(returning, table)
630
    
631
    # Build query
632
    first_line = 'INSERT INTO '+table.to_str(db)
633
    query = first_line
634
    if cols != None: query += '\n('+', '.join(cols)+')'
635
    query += '\n'+select_query
636
    
637
    if returning != None:
638
        query += '\nRETURNING '+sql_gen.to_name_only_col(returning).to_str(db)
639
    
640
    if embeddable:
641
        assert returning != None
642
        
643
        # Create function
644
        function_name = sql_gen.clean_name(first_line)
645
        return_type = 'SETOF '+returning.to_str(db)+'%TYPE'
646
        while True:
647
            try:
648
                function = db.TempFunction(function_name)
649
                
650
                function_query = '''\
651
CREATE FUNCTION '''+function.to_str(db)+'''()
652
RETURNS '''+return_type+'''
653
LANGUAGE sql
654
AS $$
655
'''+query+''';
656
$$;
657
'''
658
                run_query(db, function_query, recover=True, cacheable=True,
659
                    log_ignore_excs=(DuplicateException,))
660
                break # this version was successful
661
            except DuplicateException, e:
662
                function_name = next_version(function_name)
663
                # try again with next version of name
664
        
665
        # Return query that uses function
666
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(function),
667
            [returning]) # AS clause requires function alias
668
        return mk_select(db, func_table, start=0, order_by=None)
669
    
670
    return query
671

    
672
def insert_select(db, *args, **kw_args):
673
    '''For params, see mk_insert_select() and run_query_into()
674
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
675
        values in
676
    '''
677
    into = kw_args.pop('into', None)
678
    if into != None: kw_args['embeddable'] = True
679
    recover = kw_args.pop('recover', None)
680
    cacheable = kw_args.pop('cacheable', True)
681
    log_level = kw_args.pop('log_level', 2)
682
    
683
    return run_query_into(db, mk_insert_select(db, *args, **kw_args), into,
684
        recover=recover, cacheable=cacheable, log_level=log_level)
685

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

    
688
def insert(db, table, row, *args, **kw_args):
689
    '''For params, see insert_select()'''
690
    if lists.is_seq(row): cols = None
691
    else:
692
        cols = row.keys()
693
        row = row.values()
694
    row = list(row) # ensure that "== []" works
695
    
696
    if row == []: query = None
697
    else: query = sql_gen.Values(row).to_str(db)
698
    
699
    return insert_select(db, table, cols, query, *args, **kw_args)
700

    
701
def mk_update(db, table, changes=None, cond=None):
702
    '''
703
    @param changes [(col, new_value),...]
704
        * container can be any iterable type
705
        * col: sql_gen.Code|str (for col name)
706
        * new_value: sql_gen.Code|literal value
707
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
708
    @return str query
709
    '''
710
    query = 'UPDATE '+sql_gen.as_Table(table).to_str(db)+'\nSET\n'
711
    query += ',\n'.join((sql_gen.to_name_only_col(col, table).to_str(db)+' = '
712
        +sql_gen.as_Value(new_value).to_str(db) for col, new_value in changes))
713
    if cond != None: query += '\nWHERE\n'+cond.to_str(db)
714
    
715
    return query
716

    
717
def update(db, *args, **kw_args):
718
    '''For params, see mk_update() and run_query()'''
719
    recover = kw_args.pop('recover', None)
720
    
721
    return run_query(db, mk_update(db, *args, **kw_args), recover)
722

    
723
def last_insert_id(db):
724
    module = util.root_module(db.db)
725
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
726
    elif module == 'MySQLdb': return db.insert_id()
727
    else: return None
728

    
729
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
730
    '''Creates a mapping from original column names (which may have collisions)
731
    to names that will be distinct among the columns' tables.
732
    This is meant to be used for several tables that are being joined together.
733
    @param cols The columns to combine. Duplicates will be removed.
734
    @param into The table for the new columns.
735
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
736
        columns will be included in the mapping even if they are not in cols.
737
        The tables of the provided Col objects will be changed to into, so make
738
        copies of them if you want to keep the original tables.
739
    @param as_items Whether to return a list of dict items instead of a dict
740
    @return dict(orig_col=new_col, ...)
741
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
742
        * new_col: sql_gen.Col(orig_col_name, into)
743
        * All mappings use the into table so its name can easily be
744
          changed for all columns at once
745
    '''
746
    cols = lists.uniqify(cols)
747
    
748
    items = []
749
    for col in preserve:
750
        orig_col = copy.copy(col)
751
        col.table = into
752
        items.append((orig_col, col))
753
    preserve = set(preserve)
754
    for col in cols:
755
        if col not in preserve:
756
            items.append((col, sql_gen.Col(str(col), into, col.srcs)))
757
    
758
    if not as_items: items = dict(items)
759
    return items
760

    
761
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
762
    '''For params, see mk_flatten_mapping()
763
    @return See return value of mk_flatten_mapping()
764
    '''
765
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
766
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
767
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
768
        into=into, add_indexes_=True)
769
    return dict(items)
770

    
771
def mk_track_data_error(db, errors_table, cols, value, error_code, error):
772
    assert cols != ()
773
    
774
    cols = map(sql_gen.to_name_only_col, cols)
775
    
776
    columns_cols = ['column']
777
    columns = sql_gen.NamedValues('columns', columns_cols,
778
        [[c.name] for c in cols])
779
    values_cols = ['value', 'error_code', 'error']
780
    values = sql_gen.NamedValues('values', values_cols,
781
        [value, error_code, error])
782
    
783
    select_cols = columns_cols+values_cols
784
    name_only_cols = map(sql_gen.to_name_only_col, select_cols)
785
    errors_table = sql_gen.NamedTable('errors', errors_table)
786
    joins = [columns, sql_gen.Join(values, type_='CROSS'),
787
        sql_gen.Join(errors_table, dict(zip(name_only_cols, select_cols)),
788
        sql_gen.filter_out)]
789
    
790
    return mk_insert_select(db, errors_table, name_only_cols,
791
        mk_select(db, joins, select_cols, order_by=None))
792

    
793
def track_data_error(db, errors_table, cols, *args, **kw_args):
794
    '''
795
    @param errors_table If None, does nothing.
796
    '''
797
    if errors_table == None or cols == (): return
798
    run_query(db, mk_track_data_error(db, errors_table, cols, *args, **kw_args),
799
        cacheable=True, log_level=4)
800

    
801
def cast(db, type_, col, errors_table=None):
802
    '''Casts an (unrenamed) column or value.
803
    If errors_table set and col has srcs, saves errors in errors_table (using
804
    col's srcs attr as the source columns) and converts errors to warnings.
805
    @param col str|sql_gen.Col|sql_gen.Literal
806
    @param errors_table None|sql_gen.Table|str
807
    '''
808
    col = sql_gen.as_Col(col)
809
    save_errors = (errors_table != None and isinstance(col, sql_gen.Col)
810
        and col.srcs != ())
811
    if not save_errors: return sql_gen.Cast(type_, col) # can't save errors
812
    
813
    assert not isinstance(col, sql_gen.NamedCol)
814
    
815
    errors_table = sql_gen.as_Table(errors_table)
816
    srcs = map(sql_gen.to_name_only_col, col.srcs)
817
    function_name = str(sql_gen.FunctionCall(type_, *srcs))
818
    function = db.TempFunction(function_name)
819
    
820
    while True:
821
        # Create function definition
822
        query = '''\
823
CREATE FUNCTION '''+function.to_str(db)+'''(value text)
824
RETURNS '''+type_+'''
825
LANGUAGE plpgsql
826
STRICT
827
AS $$
828
BEGIN
829
    /* The explicit cast to the return type is needed to make the cast happen
830
    inside the try block. (Implicit casts to the return type happen at the end
831
    of the function, outside any block.) */
832
    RETURN value::'''+type_+''';
833
EXCEPTION
834
    WHEN data_exception THEN
835
        -- Save error in errors table.
836
        -- Insert the value and error for *each* source column.
837
'''+mk_track_data_error(db, errors_table, srcs,
838
    *map(sql_gen.CustomCode, ['value', 'SQLSTATE', 'SQLERRM']))+''';
839
        
840
        RAISE WARNING '%', SQLERRM;
841
        RETURN NULL;
842
END;
843
$$;
844
'''
845
        
846
        # Create function
847
        try:
848
            run_query(db, query, recover=True, cacheable=True,
849
                log_ignore_excs=(DuplicateException,))
850
            break # successful
851
        except DuplicateException:
852
            function.name = next_version(function.name)
853
            # try again with next version of name
854
    
855
    return sql_gen.FunctionCall(function, col)
856

    
857
##### Database structure queries
858

    
859
def table_row_count(db, table, recover=None):
860
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
861
        order_by=None, start=0), recover=recover, log_level=3))
862

    
863
def table_cols(db, table, recover=None):
864
    return list(col_names(select(db, table, limit=0, order_by=None,
865
        recover=recover, log_level=4)))
866

    
867
def pkey(db, table, recover=None):
868
    '''Assumed to be first column in table'''
869
    return table_cols(db, table, recover)[0]
870

    
871
not_null_col = 'not_null_col'
872

    
873
def table_not_null_col(db, table, recover=None):
874
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
875
    if not_null_col in table_cols(db, table, recover): return not_null_col
876
    else: return pkey(db, table, recover)
877

    
878
def index_cols(db, table, index):
879
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
880
    automatically created. When you don't know whether something is a UNIQUE
881
    constraint or a UNIQUE index, use this function.'''
882
    module = util.root_module(db.db)
883
    if module == 'psycopg2':
884
        return list(values(run_query(db, '''\
885
SELECT attname
886
FROM
887
(
888
        SELECT attnum, attname
889
        FROM pg_index
890
        JOIN pg_class index ON index.oid = indexrelid
891
        JOIN pg_class table_ ON table_.oid = indrelid
892
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
893
        WHERE
894
            table_.relname = '''+db.esc_value(table)+'''
895
            AND index.relname = '''+db.esc_value(index)+'''
896
    UNION
897
        SELECT attnum, attname
898
        FROM
899
        (
900
            SELECT
901
                indrelid
902
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
903
                    AS indkey
904
            FROM pg_index
905
            JOIN pg_class index ON index.oid = indexrelid
906
            JOIN pg_class table_ ON table_.oid = indrelid
907
            WHERE
908
                table_.relname = '''+db.esc_value(table)+'''
909
                AND index.relname = '''+db.esc_value(index)+'''
910
        ) s
911
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
912
) s
913
ORDER BY attnum
914
'''
915
            , cacheable=True, log_level=4)))
916
    else: raise NotImplementedError("Can't list index columns for "+module+
917
        ' database')
918

    
919
def constraint_cols(db, table, constraint):
920
    module = util.root_module(db.db)
921
    if module == 'psycopg2':
922
        return list(values(run_query(db, '''\
923
SELECT attname
924
FROM pg_constraint
925
JOIN pg_class ON pg_class.oid = conrelid
926
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
927
WHERE
928
    relname = '''+db.esc_value(table)+'''
929
    AND conname = '''+db.esc_value(constraint)+'''
930
ORDER BY attnum
931
'''
932
            )))
933
    else: raise NotImplementedError("Can't list constraint columns for "+module+
934
        ' database')
935

    
936
row_num_col = '_row_num'
937

    
938
def add_index(db, exprs, table=None, unique=False, ensure_not_null_=True):
939
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
940
    Currently, only function calls are supported as expressions.
941
    @param ensure_not_null_ If set, translates NULL values to sentinel values.
942
        This allows indexes to be used for comparisons where NULLs are equal.
943
    '''
944
    exprs = lists.mk_seq(exprs)
945
    
946
    # Parse exprs
947
    old_exprs = exprs[:]
948
    exprs = []
949
    cols = []
950
    for i, expr in enumerate(old_exprs):
951
        expr = copy.deepcopy(expr) # don't modify input!
952
        expr = sql_gen.as_Col(expr, table)
953
        
954
        # Handle nullable columns
955
        if ensure_not_null_:
956
            try: expr = ensure_not_null(db, expr)
957
            except KeyError: pass # unknown type, so just create plain index
958
        
959
        # Extract col
960
        if isinstance(expr, sql_gen.FunctionCall):
961
            col = expr.args[0]
962
            expr = sql_gen.Expr(expr)
963
        else: col = expr
964
        assert isinstance(col, sql_gen.Col)
965
        
966
        # Extract table
967
        if table == None:
968
            assert sql_gen.is_table_col(col)
969
            table = col.table
970
        
971
        col.table = None
972
        
973
        exprs.append(expr)
974
        cols.append(col)
975
    
976
    table = sql_gen.as_Table(table)
977
    index = sql_gen.Table(str(sql_gen.Col(','.join(map(str, cols)), table)))
978
    
979
    str_ = 'CREATE'
980
    if unique: str_ += ' UNIQUE'
981
    str_ += ' INDEX '+index.to_str(db)+' ON '+table.to_str(db)+' ('+(
982
        ', '.join((v.to_str(db) for v in exprs)))+')'
983
    
984
    try: run_query(db, str_, recover=True, cacheable=True, log_level=3)
985
    except DuplicateException: pass # index already existed
986

    
987
def add_pkey(db, table, cols=None, recover=None):
988
    '''Adds a primary key.
989
    @param cols [sql_gen.Col,...] The columns in the primary key.
990
        Defaults to the first column in the table.
991
    @pre The table must not already have a primary key.
992
    '''
993
    table = sql_gen.as_Table(table)
994
    if cols == None: cols = [pkey(db, table, recover)]
995
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
996
    
997
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
998
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
999
        log_ignore_excs=(DuplicateException,))
1000

    
1001
def add_not_null(db, col):
1002
    table = col.table
1003
    col = sql_gen.to_name_only_col(col)
1004
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1005
        +col.to_str(db)+' SET NOT NULL', cacheable=True)
1006

    
1007
def add_index_col(db, col, suffix, expr, nullable=True):
1008
    if col.index_col != None: return # already has index col
1009
    
1010
    new_col = sql_gen.suffixed_col(col, suffix)
1011
    new_typed_col = sql_gen.TypedCol(new_col.name, db.col_info(col).type)
1012
    
1013
    add_col(db, col.table, new_typed_col)
1014
    update(db, col.table, [(new_col, expr)])
1015
    if not nullable: add_not_null(db, new_col)
1016
    add_index(db, new_col)
1017
    
1018
    col.index_col = new_col
1019

    
1020
def ensure_not_null(db, col):
1021
    '''For params, see sql_gen.ensure_not_null()'''
1022
    expr = sql_gen.ensure_not_null(db, col)
1023
    
1024
    # If nullable column in a temp table, add separate column instead
1025
    if sql_gen.is_temp_col(col) and isinstance(expr, sql_gen.EnsureNotNull):
1026
        add_index_col(db, col, '::NOT NULL', expr, nullable=False)
1027
        expr = col.index_col
1028
    
1029
    return expr
1030

    
1031
already_indexed = object() # tells add_indexes() the pkey has already been added
1032

    
1033
def add_indexes(db, table, has_pkey=True):
1034
    '''Adds an index on all columns in a table.
1035
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
1036
        index should be added on the first column.
1037
        * If already_indexed, the pkey is assumed to have already been added
1038
    '''
1039
    cols = table_cols(db, table)
1040
    if has_pkey:
1041
        if has_pkey is not already_indexed: add_pkey(db, table)
1042
        cols = cols[1:]
1043
    for col in cols: add_index(db, col, table)
1044

    
1045
def add_col(db, table, col, **kw_args):
1046
    assert isinstance(col, sql_gen.TypedCol)
1047
    try: run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '
1048
        +col.to_str(db), recover=True, cacheable=True, **kw_args)
1049
    except DuplicateException: pass # column already existed
1050

    
1051
row_num_typed_col = sql_gen.TypedCol(row_num_col, 'serial', nullable=False,
1052
    constraints='PRIMARY KEY')
1053

    
1054
def add_row_num(db, table):
1055
    '''Adds a row number column to a table. Its name is in row_num_col. It will
1056
    be the primary key.'''
1057
    add_col(db, table, row_num_typed_col, log_level=3)
1058

    
1059
def cast_temp_col(db, type_, col, errors_table=None):
1060
    '''Like cast(), but creates a new column with the cast values if the input
1061
    is a column.
1062
    @return The new column or cast value
1063
    '''
1064
    def cast_(col): return cast(db, type_, col, errors_table)
1065
    
1066
    try: col = sql_gen.underlying_col(col)
1067
    except sql_gen.NoUnderlyingTableException: return sql_gen.wrap(cast_, col)
1068
    
1069
    table = col.table
1070
    new_col = sql_gen.Col(sql_gen.concat(col.name, '::'+type_), table, col.srcs)
1071
    expr = cast_(col)
1072
    add_col(db, table, sql_gen.TypedCol(new_col.name, type_))
1073
    update(db, table, [(new_col, expr)])
1074
    add_index(db, new_col)
1075
    
1076
    return new_col
1077

    
1078
def drop_table(db, table):
1079
    table = sql_gen.as_Table(table)
1080
    return run_query(db, 'DROP TABLE IF EXISTS '+table.to_str(db)+' CASCADE')
1081

    
1082
def create_table(db, table, cols, has_pkey=True, col_indexes=True):
1083
    '''Creates a table.
1084
    @param cols [sql_gen.TypedCol,...] The column names and types
1085
    @param has_pkey If set, the first column becomes the primary key.
1086
    @param col_indexes bool|[ref]
1087
        * If True, indexes will be added on all non-pkey columns.
1088
        * If a list reference, [0] will be set to a function to do this.
1089
          This can be used to delay index creation until the table is populated.
1090
    '''
1091
    table = sql_gen.as_Table(table)
1092
    
1093
    if has_pkey:
1094
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1095
        pkey.constraints = 'PRIMARY KEY'
1096
    
1097
    str_ = 'CREATE TABLE '+table.to_str(db)+' (\n'
1098
    str_ += '\n, '.join(v.to_str(db) for v in cols)
1099
    str_ += '\n);\n'
1100
    run_query(db, str_, cacheable=True, log_level=2)
1101
    
1102
    # Add indexes
1103
    if has_pkey: has_pkey = already_indexed
1104
    def add_indexes_(): add_indexes(db, table, has_pkey)
1105
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1106
    elif col_indexes: add_indexes_() # add now
1107

    
1108
def vacuum(db, table):
1109
    table = sql_gen.as_Table(table)
1110
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1111
        log_level=3))
1112

    
1113
def truncate(db, table, schema='public', **kw_args):
1114
    '''For params, see run_query()'''
1115
    table = sql_gen.as_Table(table, schema)
1116
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE', **kw_args)
1117

    
1118
def empty_temp(db, tables):
1119
    if db.debug_temp: return # leave temp tables there for debugging
1120
    tables = lists.mk_seq(tables)
1121
    for table in tables: truncate(db, table, log_level=3)
1122

    
1123
def tables(db, schema_like='public', table_like='%', exact=False):
1124
    if exact: compare = '='
1125
    else: compare = 'LIKE'
1126
    
1127
    module = util.root_module(db.db)
1128
    if module == 'psycopg2':
1129
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1130
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1131
        return values(select(db, 'pg_tables', ['tablename'], conds,
1132
            order_by='tablename', log_level=4))
1133
    elif module == 'MySQLdb':
1134
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1135
            , cacheable=True, log_level=4))
1136
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1137

    
1138
def table_exists(db, table):
1139
    table = sql_gen.as_Table(table)
1140
    return list(tables(db, table.schema, table.name, exact=True)) != []
1141

    
1142
def function_exists(db, function):
1143
    function = sql_gen.as_Function(function)
1144
    
1145
    info_table = sql_gen.Table('routines', 'information_schema')
1146
    conds = [('routine_name', function.name)]
1147
    schema = function.schema
1148
    if schema != None: conds.append(('routine_schema', schema))
1149
    # Exclude trigger functions, since they cannot be called directly
1150
    conds.append(('data_type', sql_gen.CompareCond('trigger', '!=')))
1151
    
1152
    return list(values(select(db, info_table, ['routine_name'], conds,
1153
        order_by='routine_schema', limit=1, log_level=4))) != []
1154
        # TODO: order_by search_path schema order
1155

    
1156
def errors_table(db, table, if_exists=True):
1157
    '''
1158
    @param if_exists If set, returns None if the errors table doesn't exist
1159
    @return None|sql_gen.Table
1160
    '''
1161
    table = sql_gen.as_Table(table)
1162
    if table.srcs != (): table = table.srcs[0]
1163
    
1164
    errors_table = sql_gen.suffixed_table(table, '.errors')
1165
    if if_exists and not table_exists(db, errors_table): return None
1166
    return errors_table
1167

    
1168
##### Database management
1169

    
1170
def empty_db(db, schema='public', **kw_args):
1171
    '''For kw_args, see tables()'''
1172
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1173

    
1174
##### Heuristic queries
1175

    
1176
def put(db, table, row, pkey_=None, row_ct_ref=None):
1177
    '''Recovers from errors.
1178
    Only works under PostgreSQL (uses INSERT RETURNING).
1179
    '''
1180
    row = sql_gen.ColDict(db, table, row)
1181
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
1182
    
1183
    try:
1184
        cur = insert(db, table, row, pkey_, recover=True)
1185
        if row_ct_ref != None and cur.rowcount >= 0:
1186
            row_ct_ref[0] += cur.rowcount
1187
        return value(cur)
1188
    except DuplicateKeyException, e:
1189
        row = sql_gen.ColDict(db, table,
1190
            util.dict_subset_right_join(row, e.cols))
1191
        return value(select(db, table, [pkey_], row, recover=True))
1192

    
1193
def get(db, table, row, pkey, row_ct_ref=None, create=False):
1194
    '''Recovers from errors'''
1195
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
1196
    except StopIteration:
1197
        if not create: raise
1198
        return put(db, table, row, pkey, row_ct_ref) # insert new row
1199

    
1200
def is_func_result(col):
1201
    return col.table.name.find('(') >= 0 and col.name == 'result'
1202

    
1203
def into_table_name(out_table, in_tables0, mapping, is_func):
1204
    def in_col_str(in_col):
1205
        in_col = sql_gen.remove_col_rename(in_col)
1206
        if isinstance(in_col, sql_gen.Col):
1207
            table = in_col.table
1208
            if table == in_tables0:
1209
                in_col = sql_gen.to_name_only_col(in_col)
1210
            elif is_func_result(in_col): in_col = table # omit col name
1211
        return str(in_col)
1212
    
1213
    str_ = str(out_table)
1214
    if is_func:
1215
        str_ += '('
1216
        
1217
        try: value_in_col = mapping['value']
1218
        except KeyError:
1219
            str_ += ', '.join((str(k)+'='+in_col_str(v)
1220
                for k, v in mapping.iteritems()))
1221
        else: str_ += in_col_str(value_in_col)
1222
        
1223
        str_ += ')'
1224
    else:
1225
        out_col = 'rank'
1226
        try: in_col = mapping[out_col]
1227
        except KeyError: str_ += '_pkeys'
1228
        else: # has a rank column, so hierarchical
1229
            str_ += '['+str(out_col)+'='+in_col_str(in_col)+']'
1230
    return str_
1231

    
1232
def put_table(db, out_table, in_tables, mapping, row_ct_ref=None, into=None,
1233
    default=None, is_func=False, on_error=exc.raise_):
1234
    '''Recovers from errors.
1235
    Only works under PostgreSQL (uses INSERT RETURNING).
1236
    IMPORTANT: Must be run at the *beginning* of a transaction.
1237
    @param in_tables The main input table to select from, followed by a list of
1238
        tables to join with it using the main input table's pkey
1239
    @param mapping dict(out_table_col=in_table_col, ...)
1240
        * out_table_col: str (*not* sql_gen.Col)
1241
        * in_table_col: sql_gen.Col|literal-value
1242
    @param into The table to contain the output and input pkeys.
1243
        Defaults to `out_table.name+'_pkeys'`.
1244
    @param default The *output* column to use as the pkey for missing rows.
1245
        If this output column does not exist in the mapping, uses None.
1246
    @param is_func Whether out_table is the name of a SQL function, not a table
1247
    @return sql_gen.Col Where the output pkeys are made available
1248
    '''
1249
    out_table = sql_gen.as_Table(out_table)
1250
    
1251
    def log_debug(msg): db.log_debug(msg, level=1.5)
1252
    def col_ustr(str_):
1253
        return strings.repr_no_u(sql_gen.remove_col_rename(str_))
1254
    
1255
    log_debug('********** New iteration **********')
1256
    log_debug('Inserting these input columns into '+strings.as_tt(
1257
        out_table.to_str(db))+':\n'+strings.as_table(mapping, ustr=col_ustr))
1258
    
1259
    is_function = function_exists(db, out_table)
1260
    
1261
    if is_function: out_pkey = 'result'
1262
    else: out_pkey = pkey(db, out_table, recover=True)
1263
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
1264
    
1265
    if mapping == {}: # need at least one column for INSERT SELECT
1266
        mapping = {out_pkey: None} # ColDict will replace with default value
1267
    
1268
    # Create input joins from list of input tables
1269
    in_tables_ = in_tables[:] # don't modify input!
1270
    in_tables0 = in_tables_.pop(0) # first table is separate
1271
    errors_table_ = errors_table(db, in_tables0)
1272
    in_pkey = pkey(db, in_tables0, recover=True)
1273
    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
1274
    input_joins = [in_tables0]+[sql_gen.Join(v,
1275
        {in_pkey: sql_gen.join_same_not_null}) for v in in_tables_]
1276
    
1277
    if into == None:
1278
        into = into_table_name(out_table, in_tables0, mapping, is_func)
1279
    into = sql_gen.as_Table(into)
1280
    
1281
    # Set column sources
1282
    in_cols = filter(sql_gen.is_table_col, mapping.values())
1283
    for col in in_cols:
1284
        if col.table == in_tables0: col.set_srcs(sql_gen.src_self)
1285
    
1286
    log_debug('Joining together input tables into temp table')
1287
    # Place in new table for speed and so don't modify input if values edited
1288
    in_table = sql_gen.Table('in')
1289
    mapping = dicts.join(mapping, flatten(db, in_table, input_joins, in_cols,
1290
        preserve=[in_pkey_col], start=0))
1291
    input_joins = [in_table]
1292
    db.log_debug('Temp table: '+strings.as_tt(in_table.to_str(db)), level=2)
1293
    
1294
    mapping = sql_gen.ColDict(db, out_table, mapping)
1295
        # after applying dicts.join() because that returns a plain dict
1296
    
1297
    # Resolve default value column
1298
    if default != None:
1299
        try: default = mapping[default]
1300
        except KeyError:
1301
            db.log_debug('Default value column '
1302
                +strings.as_tt(strings.repr_no_u(default))
1303
                +' does not exist in mapping, falling back to None', level=2.1)
1304
            default = None
1305
    
1306
    pkeys_names = [in_pkey, out_pkey]
1307
    pkeys_cols = [in_pkey_col, out_pkey_col]
1308
    
1309
    pkeys_table_exists_ref = [False]
1310
    def insert_into_pkeys(joins, cols, distinct=False):
1311
        kw_args = {}
1312
        if distinct: kw_args.update(dict(distinct_on=[in_pkey_col]))
1313
        query = mk_select(db, joins, cols, order_by=None, start=0, **kw_args)
1314
        
1315
        if pkeys_table_exists_ref[0]:
1316
            insert_select(db, into, pkeys_names, query)
1317
        else:
1318
            run_query_into(db, query, into=into)
1319
            pkeys_table_exists_ref[0] = True
1320
    
1321
    limit_ref = [None]
1322
    conds = set()
1323
    distinct_on = sql_gen.ColDict(db, out_table)
1324
    def mk_main_select(joins, cols):
1325
        distinct_on_cols = [c.to_Col() for c in distinct_on.values()]
1326
        return mk_select(db, joins, cols, conds, distinct_on_cols,
1327
            limit=limit_ref[0], start=0)
1328
    
1329
    exc_strs = set()
1330
    def log_exc(e):
1331
        e_str = exc.str_(e, first_line_only=True)
1332
        log_debug('Caught exception: '+e_str)
1333
        assert e_str not in exc_strs # avoid infinite loops
1334
        exc_strs.add(e_str)
1335
    
1336
    def remove_all_rows():
1337
        log_debug('Ignoring all rows')
1338
        limit_ref[0] = 0 # just create an empty pkeys table
1339
    
1340
    def ignore(in_col, value, e):
1341
        track_data_error(db, errors_table_, in_col.srcs, value, e.cause.pgcode,
1342
            e.cause.pgerror)
1343
        log_debug('Ignoring rows with '+strings.as_tt(repr(in_col))+' = '
1344
            +strings.as_tt(repr(value)))
1345
    
1346
    def remove_rows(in_col, value, e):
1347
        ignore(in_col, value, e)
1348
        cond = (in_col, sql_gen.CompareCond(value, '!='))
1349
        assert cond not in conds # avoid infinite loops
1350
        conds.add(cond)
1351
    
1352
    def invalid2null(in_col, value, e):
1353
        ignore(in_col, value, e)
1354
        update(db, in_table, [(in_col, None)],
1355
            sql_gen.ColValueCond(in_col, value))
1356
    
1357
    def insert_pkeys_table(which):
1358
        return sql_gen.Table(sql_gen.concat(in_table.name,
1359
            '_insert_'+which+'_pkeys'))
1360
    insert_out_pkeys = insert_pkeys_table('out')
1361
    insert_in_pkeys = insert_pkeys_table('in')
1362
    
1363
    # Do inserts and selects
1364
    join_cols = sql_gen.ColDict(db, out_table)
1365
    while True:
1366
        if limit_ref[0] == 0: # special case
1367
            log_debug('Creating an empty pkeys table')
1368
            cur = run_query_into(db, mk_select(db, out_table, [out_pkey],
1369
                limit=limit_ref[0]), into=insert_out_pkeys)
1370
            break # don't do main case
1371
        
1372
        has_joins = join_cols != {}
1373
        
1374
        log_debug('Trying to insert new rows')
1375
        
1376
        # Prepare to insert new rows
1377
        insert_joins = input_joins[:] # don't modify original!
1378
        insert_args = dict(recover=True, cacheable=False)
1379
        if has_joins:
1380
            insert_joins.append(sql_gen.Join(out_table, join_cols,
1381
                sql_gen.filter_out))
1382
        else:
1383
            insert_args.update(dict(returning=out_pkey, into=insert_out_pkeys))
1384
        main_select = mk_main_select(insert_joins, mapping.values())
1385
        
1386
        def main_insert():
1387
            if is_function:
1388
                log_debug('Calling function on input rows')
1389
                args = dict(((k.name, v) for k, v in mapping.iteritems()))
1390
                func_call = sql_gen.NamedCol(out_pkey,
1391
                    sql_gen.FunctionCall(out_table, **args))
1392
                insert_into_pkeys(input_joins, [in_pkey_col, func_call])
1393
                return None
1394
            else:
1395
                lock_table(db, out_table, 'EXCLUSIVE')
1396
                return insert_select(db, out_table, mapping.keys(), main_select,
1397
                    **insert_args)
1398
        
1399
        try:
1400
            cur = with_savepoint(db, main_insert)
1401
            break # insert successful
1402
        except MissingCastException, e:
1403
            log_exc(e)
1404
            
1405
            out_col = e.col
1406
            type_ = e.type
1407
            
1408
            log_debug('Casting '+strings.as_tt(out_col)+' input to '
1409
                +strings.as_tt(type_))
1410
            mapping[out_col] = cast_temp_col(db, type_, mapping[out_col],
1411
                errors_table_)
1412
        except DuplicateKeyException, e:
1413
            log_exc(e)
1414
            
1415
            old_join_cols = join_cols.copy()
1416
            distinct_on.update(util.dict_subset(mapping, e.cols))
1417
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
1418
            log_debug('Ignoring existing rows, comparing on these columns:\n'
1419
                +strings.as_inline_table(join_cols, ustr=col_ustr))
1420
            assert join_cols != old_join_cols # avoid infinite loops
1421
        except NullValueException, e:
1422
            log_exc(e)
1423
            
1424
            out_col, = e.cols
1425
            try: in_col = mapping[out_col]
1426
            except KeyError:
1427
                log_debug('Missing mapping for NOT NULL column '+out_col)
1428
                remove_all_rows()
1429
            else: remove_rows(in_col, None, e)
1430
        except FunctionValueException, e:
1431
            log_exc(e)
1432
            
1433
            func_name = e.name
1434
            value = e.value
1435
            for out_col, in_col in mapping.iteritems():
1436
                in_col = sql_gen.unwrap_func_call(in_col, func_name)
1437
                invalid2null(in_col, value, e)
1438
        except DatabaseErrors, e:
1439
            log_exc(e)
1440
            
1441
            log_debug('No handler for exception')
1442
            on_error(e)
1443
            remove_all_rows()
1444
        # after exception handled, rerun loop with additional constraints
1445
    
1446
    if cur != None and row_ct_ref != None and cur.rowcount >= 0:
1447
        row_ct_ref[0] += cur.rowcount
1448
    
1449
    if is_function: pass # pkeys table already created
1450
    elif has_joins:
1451
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
1452
        log_debug('Getting output table pkeys of existing/inserted rows')
1453
        insert_into_pkeys(select_joins, pkeys_cols, distinct=True)
1454
    else:
1455
        add_row_num(db, insert_out_pkeys) # for joining with input pkeys
1456
        
1457
        log_debug('Getting input table pkeys of inserted rows')
1458
        run_query_into(db, mk_main_select(input_joins, [in_pkey]),
1459
            into=insert_in_pkeys)
1460
        add_row_num(db, insert_in_pkeys) # for joining with output pkeys
1461
        
1462
        assert table_row_count(db, insert_out_pkeys) == table_row_count(db,
1463
            insert_in_pkeys)
1464
        
1465
        log_debug('Combining output and input pkeys in inserted order')
1466
        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,
1467
            {row_num_col: sql_gen.join_same_not_null})]
1468
        insert_into_pkeys(pkey_joins, pkeys_names)
1469
        
1470
        empty_temp(db, [insert_out_pkeys, insert_in_pkeys])
1471
    
1472
    db.log_debug('Adding pkey on pkeys table to enable fast joins', level=2.5)
1473
    add_pkey(db, into)
1474
    
1475
    log_debug('Setting pkeys of missing rows to '+strings.as_tt(repr(default)))
1476
    missing_rows_joins = input_joins+[sql_gen.Join(into,
1477
        {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
1478
        # must use join_same_not_null or query will take forever
1479
    insert_into_pkeys(missing_rows_joins,
1480
        [in_pkey_col, sql_gen.NamedCol(out_pkey, default)])
1481
    
1482
    assert table_row_count(db, into) == table_row_count(db, in_table)
1483
    
1484
    empty_temp(db, in_table)
1485
    
1486
    srcs = []
1487
    if is_func: srcs = sql_gen.cols_srcs(in_cols)
1488
    return sql_gen.Col(out_pkey, into, srcs)
1489

    
1490
##### Data cleanup
1491

    
1492
def cleanup_table(db, table, cols):
1493
    table = sql_gen.as_Table(table)
1494
    cols = map(sql_gen.as_Col, cols)
1495
    
1496
    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '
1497
        +db.esc_value(r'\N')+')')
1498
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))
1499
        for v in cols]
1500
    
1501
    update(db, table, changes)
(24-24/36)