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 ConstraintException(DbException):
50
    def __init__(self, name, cols, cause=None):
51
        DbException.__init__(self, 'Violated '+strings.as_tt(name)
52
            +' constraint on columns: '+strings.as_tt(', '.join(cols)), cause)
53
        self.name = name
54
        self.cols = cols
55

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

    
63
class NameException(DbException): pass
64

    
65
class DuplicateKeyException(ConstraintException): pass
66

    
67
class NullValueException(ConstraintException): pass
68

    
69
class FunctionValueException(ExceptionWithNameValue): pass
70

    
71
class DuplicateTableException(ExceptionWithName): pass
72

    
73
class DuplicateFunctionException(ExceptionWithName): pass
74

    
75
class EmptyRowException(DbException): pass
76

    
77
##### Warnings
78

    
79
class DbWarning(UserWarning): pass
80

    
81
##### Result retrieval
82

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

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

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

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

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

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

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

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

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

    
108
##### Escaping
109

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

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

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

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

    
128
##### Database connections
129

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

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

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

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

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

    
148
log_debug_none = lambda msg, level=2: None
149

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

    
398
connect = DbConn
399

    
400
##### Recoverable querying
401

    
402
def with_savepoint(db, func): return db.with_savepoint(func)
403

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

    
465
##### Basic queries
466

    
467
def next_version(name):
468
    version = 1 # first existing name was version 0
469
    match = re.match(r'^(.*)#(\d+)$', name)
470
    if match:
471
        name, version = match.groups()
472
        version = int(version)+1
473
    return sql_gen.add_suffix(name, '#'+str(version))
474

    
475
def lock_table(db, table, mode):
476
    table = sql_gen.as_Table(table)
477
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
478

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

    
512
order_by_pkey = object() # tells mk_select() to order by the pkey
513

    
514
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
515

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

    
597
def select(db, *args, **kw_args):
598
    '''For params, see mk_select() and run_query()'''
599
    recover = kw_args.pop('recover', None)
600
    cacheable = kw_args.pop('cacheable', True)
601
    log_level = kw_args.pop('log_level', 2)
602
    
603
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
604
        log_level=log_level)
605

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

    
662
def insert_select(db, *args, **kw_args):
663
    '''For params, see mk_insert_select() and run_query_into()
664
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
665
        values in
666
    '''
667
    into = kw_args.pop('into', None)
668
    if into != None: kw_args['embeddable'] = True
669
    recover = kw_args.pop('recover', None)
670
    cacheable = kw_args.pop('cacheable', True)
671
    log_level = kw_args.pop('log_level', 2)
672
    
673
    return run_query_into(db, mk_insert_select(db, *args, **kw_args), into,
674
        recover=recover, cacheable=cacheable, log_level=log_level)
675

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

    
678
def insert(db, table, row, *args, **kw_args):
679
    '''For params, see insert_select()'''
680
    if lists.is_seq(row): cols = None
681
    else:
682
        cols = row.keys()
683
        row = row.values()
684
    row = list(row) # ensure that "== []" works
685
    
686
    if row == []: query = None
687
    else: query = sql_gen.Values(row).to_str(db)
688
    
689
    return insert_select(db, table, cols, query, *args, **kw_args)
690

    
691
def mk_update(db, table, changes=None, cond=None):
692
    '''
693
    @param changes [(col, new_value),...]
694
        * container can be any iterable type
695
        * col: sql_gen.Code|str (for col name)
696
        * new_value: sql_gen.Code|literal value
697
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
698
    @return str query
699
    '''
700
    query = 'UPDATE '+sql_gen.as_Table(table).to_str(db)+'\nSET\n'
701
    query += ',\n'.join((sql_gen.to_name_only_col(col, table).to_str(db)+' = '
702
        +sql_gen.as_Value(new_value).to_str(db) for col, new_value in changes))
703
    if cond != None: query += '\nWHERE\n'+cond.to_str(db)
704
    
705
    return query
706

    
707
def update(db, *args, **kw_args):
708
    '''For params, see mk_update() and run_query()'''
709
    recover = kw_args.pop('recover', None)
710
    
711
    return run_query(db, mk_update(db, *args, **kw_args), recover)
712

    
713
def last_insert_id(db):
714
    module = util.root_module(db.db)
715
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
716
    elif module == 'MySQLdb': return db.insert_id()
717
    else: return None
718

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

    
751
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
752
    '''For params, see mk_flatten_mapping()
753
    @return See return value of mk_flatten_mapping()
754
    '''
755
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
756
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
757
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
758
        into=into, add_indexes_=True)
759
    return dict(items)
760

    
761
def mk_track_data_error(db, errors_table, cols, value, error_code, error):
762
    assert cols != ()
763
    
764
    cols = map(sql_gen.to_name_only_col, cols)
765
    
766
    columns_cols = ['column']
767
    columns = sql_gen.NamedValues('columns', columns_cols,
768
        [[c.name] for c in cols])
769
    values_cols = ['value', 'error_code', 'error']
770
    values = sql_gen.NamedValues('values', values_cols,
771
        [value, error_code, error])
772
    
773
    select_cols = columns_cols+values_cols
774
    name_only_cols = map(sql_gen.to_name_only_col, select_cols)
775
    errors_table = sql_gen.NamedTable('errors', errors_table)
776
    joins = [columns, sql_gen.Join(values, type_='CROSS'),
777
        sql_gen.Join(errors_table, dict(zip(name_only_cols, select_cols)),
778
        sql_gen.filter_out)]
779
    
780
    return mk_insert_select(db, errors_table, name_only_cols,
781
        mk_select(db, joins, select_cols, order_by=None))
782

    
783
def track_data_error(db, errors_table, cols, *args, **kw_args):
784
    '''
785
    @param errors_table If None, does nothing.
786
    '''
787
    if errors_table == None or cols == (): return
788
    run_query(db, mk_track_data_error(db, errors_table, cols, *args, **kw_args),
789
        cacheable=True, log_level=4)
790

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

    
848
##### Database structure queries
849

    
850
def table_row_count(db, table, recover=None):
851
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
852
        order_by=None, start=0), recover=recover, log_level=3))
853

    
854
def table_cols(db, table, recover=None):
855
    return list(col_names(select(db, table, limit=0, order_by=None,
856
        recover=recover, log_level=4)))
857

    
858
def pkey(db, table, recover=None):
859
    '''Assumed to be first column in table'''
860
    return table_cols(db, table, recover)[0]
861

    
862
not_null_col = 'not_null_col'
863

    
864
def table_not_null_col(db, table, recover=None):
865
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
866
    if not_null_col in table_cols(db, table, recover): return not_null_col
867
    else: return pkey(db, table, recover)
868

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

    
910
def constraint_cols(db, table, constraint):
911
    module = util.root_module(db.db)
912
    if module == 'psycopg2':
913
        return list(values(run_query(db, '''\
914
SELECT attname
915
FROM pg_constraint
916
JOIN pg_class ON pg_class.oid = conrelid
917
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
918
WHERE
919
    relname = '''+db.esc_value(table)+'''
920
    AND conname = '''+db.esc_value(constraint)+'''
921
ORDER BY attnum
922
'''
923
            )))
924
    else: raise NotImplementedError("Can't list constraint columns for "+module+
925
        ' database')
926

    
927
row_num_col = '_row_num'
928

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

    
978
def add_pkey(db, table, cols=None, recover=None):
979
    '''Adds a primary key.
980
    @param cols [sql_gen.Col,...] The columns in the primary key.
981
        Defaults to the first column in the table.
982
    @pre The table must not already have a primary key.
983
    '''
984
    table = sql_gen.as_Table(table)
985
    if cols == None: cols = [pkey(db, table, recover)]
986
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
987
    
988
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
989
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
990
        log_ignore_excs=(DuplicateTableException,))
991

    
992
already_indexed = object() # tells add_indexes() the pkey has already been added
993

    
994
def add_indexes(db, table, has_pkey=True):
995
    '''Adds an index on all columns in a table.
996
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
997
        index should be added on the first column.
998
        * If already_indexed, the pkey is assumed to have already been added
999
    '''
1000
    cols = table_cols(db, table)
1001
    if has_pkey:
1002
        if has_pkey is not already_indexed: add_pkey(db, table)
1003
        cols = cols[1:]
1004
    for col in cols: add_index(db, col, table)
1005

    
1006
def add_row_num(db, table):
1007
    '''Adds a row number column to a table. Its name is in row_num_col. It will
1008
    be the primary key.'''
1009
    table = sql_gen.as_Table(table).to_str(db)
1010
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
1011
        +' serial NOT NULL PRIMARY KEY', log_level=3)
1012

    
1013
def drop_table(db, table):
1014
    table = sql_gen.as_Table(table)
1015
    return run_query(db, 'DROP TABLE IF EXISTS '+table.to_str(db)+' CASCADE')
1016

    
1017
def create_table(db, table, cols, has_pkey=True, col_indexes=True):
1018
    '''Creates a table.
1019
    @param cols [sql_gen.TypedCol,...] The column names and types
1020
    @param has_pkey If set, the first column becomes the primary key.
1021
    @param col_indexes bool|[ref]
1022
        * If True, indexes will be added on all non-pkey columns.
1023
        * If a list reference, [0] will be set to a function to do this.
1024
          This can be used to delay index creation until the table is populated.
1025
    '''
1026
    table = sql_gen.as_Table(table)
1027
    
1028
    if has_pkey:
1029
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1030
        pkey.constraints = 'PRIMARY KEY'
1031
    
1032
    str_ = 'CREATE TABLE '+table.to_str(db)+' (\n'
1033
    str_ += '\n, '.join(v.to_str(db) for v in cols)
1034
    str_ += '\n);\n'
1035
    run_query(db, str_, cacheable=True, log_level=2)
1036
    
1037
    # Add indexes
1038
    if has_pkey: has_pkey = already_indexed
1039
    def add_indexes_(): add_indexes(db, table, has_pkey)
1040
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1041
    elif col_indexes: add_indexes_() # add now
1042

    
1043
def vacuum(db, table):
1044
    table = sql_gen.as_Table(table)
1045
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1046
        log_level=3))
1047

    
1048
def truncate(db, table, schema='public'):
1049
    table = sql_gen.as_Table(table, schema)
1050
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE')
1051

    
1052
def tables(db, schema_like='public', table_like='%', exact=False):
1053
    if exact: compare = '='
1054
    else: compare = 'LIKE'
1055
    
1056
    module = util.root_module(db.db)
1057
    if module == 'psycopg2':
1058
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1059
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1060
        return values(select(db, 'pg_tables', ['tablename'], conds,
1061
            order_by='tablename', log_level=4))
1062
    elif module == 'MySQLdb':
1063
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1064
            , cacheable=True, log_level=4))
1065
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1066

    
1067
def table_exists(db, table):
1068
    table = sql_gen.as_Table(table)
1069
    return list(tables(db, table.schema, table.name, exact=True)) != []
1070

    
1071
def function_exists(db, function):
1072
    function = sql_gen.as_Function(function)
1073
    
1074
    info_table = sql_gen.Table('routines', 'information_schema')
1075
    conds = [('routine_name', function.name)]
1076
    schema = function.schema
1077
    if schema != None: conds.append(('routine_schema', schema))
1078
    # Exclude trigger functions, since they cannot be called directly
1079
    conds.append(('data_type', sql_gen.CompareCond('trigger', '!=')))
1080
    
1081
    return list(values(select(db, info_table, ['routine_name'], conds,
1082
        order_by='routine_schema', limit=1, log_level=4))) != []
1083
        # TODO: order_by search_path schema order
1084

    
1085
def errors_table(db, table, if_exists=True):
1086
    '''
1087
    @param if_exists If set, returns None if the errors table doesn't exist
1088
    @return None|sql_gen.Table
1089
    '''
1090
    table = sql_gen.as_Table(table)
1091
    if table.srcs != (): table = table.srcs[0]
1092
    
1093
    errors_table = sql_gen.suffixed_table(table, '.errors')
1094
    if if_exists and not table_exists(db, errors_table): return None
1095
    return errors_table
1096

    
1097
##### Database management
1098

    
1099
def empty_db(db, schema='public', **kw_args):
1100
    '''For kw_args, see tables()'''
1101
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1102

    
1103
##### Heuristic queries
1104

    
1105
def put(db, table, row, pkey_=None, row_ct_ref=None):
1106
    '''Recovers from errors.
1107
    Only works under PostgreSQL (uses INSERT RETURNING).
1108
    '''
1109
    row = sql_gen.ColDict(db, table, row)
1110
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
1111
    
1112
    try:
1113
        cur = insert(db, table, row, pkey_, recover=True)
1114
        if row_ct_ref != None and cur.rowcount >= 0:
1115
            row_ct_ref[0] += cur.rowcount
1116
        return value(cur)
1117
    except DuplicateKeyException, e:
1118
        row = sql_gen.ColDict(db, table,
1119
            util.dict_subset_right_join(row, e.cols))
1120
        return value(select(db, table, [pkey_], row, recover=True))
1121

    
1122
def get(db, table, row, pkey, row_ct_ref=None, create=False):
1123
    '''Recovers from errors'''
1124
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
1125
    except StopIteration:
1126
        if not create: raise
1127
        return put(db, table, row, pkey, row_ct_ref) # insert new row
1128

    
1129
def is_func_result(col):
1130
    return col.table.name.find('(') >= 0 and col.name == 'result'
1131

    
1132
def into_table_name(out_table, in_tables0, mapping, is_func):
1133
    def in_col_str(in_col):
1134
        in_col = sql_gen.remove_col_rename(in_col)
1135
        if isinstance(in_col, sql_gen.Col):
1136
            table = in_col.table
1137
            if table == in_tables0:
1138
                in_col = sql_gen.to_name_only_col(in_col)
1139
            elif is_func_result(in_col): in_col = table # omit col name
1140
        return str(in_col)
1141
    
1142
    str_ = str(out_table)
1143
    if is_func:
1144
        str_ += '('
1145
        
1146
        try: value_in_col = mapping['value']
1147
        except KeyError:
1148
            str_ += ', '.join((str(k)+'='+in_col_str(v)
1149
                for k, v in mapping.iteritems()))
1150
        else: str_ += in_col_str(value_in_col)
1151
        
1152
        str_ += ')'
1153
    else:
1154
        out_col = 'rank'
1155
        try: in_col = mapping[out_col]
1156
        except KeyError: str_ += '_pkeys'
1157
        else: # has a rank column, so hierarchical
1158
            str_ += '['+str(out_col)+'='+in_col_str(in_col)+']'
1159
    return str_
1160

    
1161
def put_table(db, out_table, in_tables, mapping, row_ct_ref=None, into=None,
1162
    default=None, is_func=False, on_error=exc.raise_):
1163
    '''Recovers from errors.
1164
    Only works under PostgreSQL (uses INSERT RETURNING).
1165
    IMPORTANT: Must be run at the *beginning* of a transaction.
1166
    @param in_tables The main input table to select from, followed by a list of
1167
        tables to join with it using the main input table's pkey
1168
    @param mapping dict(out_table_col=in_table_col, ...)
1169
        * out_table_col: str (*not* sql_gen.Col)
1170
        * in_table_col: sql_gen.Col|literal-value
1171
    @param into The table to contain the output and input pkeys.
1172
        Defaults to `out_table.name+'_pkeys'`.
1173
    @param default The *output* column to use as the pkey for missing rows.
1174
        If this output column does not exist in the mapping, uses None.
1175
    @param is_func Whether out_table is the name of a SQL function, not a table
1176
    @return sql_gen.Col Where the output pkeys are made available
1177
    '''
1178
    out_table = sql_gen.as_Table(out_table)
1179
    
1180
    def log_debug(msg): db.log_debug(msg, level=1.5)
1181
    def col_ustr(str_):
1182
        return strings.repr_no_u(sql_gen.remove_col_rename(str_))
1183
    
1184
    log_debug('********** New iteration **********')
1185
    log_debug('Inserting these input columns into '+strings.as_tt(
1186
        out_table.to_str(db))+':\n'+strings.as_table(mapping, ustr=col_ustr))
1187
    
1188
    log_debug('Locking output table to prevent concurrent duplicate keys')
1189
    # Must happen before any SELECT query on the table to avoid lock upgrades
1190
    lock_table(db, out_table, 'EXCLUSIVE')
1191
    
1192
    out_pkey = pkey(db, out_table, recover=True)
1193
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
1194
    
1195
    if mapping == {}: # need at least one column for INSERT SELECT
1196
        mapping = {out_pkey: None} # ColDict will replace with default value
1197
    
1198
    # Create input joins from list of input tables
1199
    in_tables_ = in_tables[:] # don't modify input!
1200
    in_tables0 = in_tables_.pop(0) # first table is separate
1201
    errors_table_ = errors_table(db, in_tables0)
1202
    in_pkey = pkey(db, in_tables0, recover=True)
1203
    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
1204
    input_joins = [in_tables0]+[sql_gen.Join(v,
1205
        {in_pkey: sql_gen.join_same_not_null}) for v in in_tables_]
1206
    
1207
    if into == None:
1208
        into = into_table_name(out_table, in_tables0, mapping, is_func)
1209
    into = sql_gen.as_Table(into)
1210
    
1211
    # Set column sources
1212
    in_cols = filter(sql_gen.is_table_col, mapping.values())
1213
    for col in in_cols:
1214
        if col.table == in_tables0: col.set_srcs(sql_gen.src_self)
1215
    
1216
    log_debug('Joining together input tables into temp table')
1217
    # Place in new table for speed and so don't modify input if values edited
1218
    in_table = sql_gen.Table('in')
1219
    mapping = dicts.join(mapping, flatten(db, in_table, input_joins, in_cols,
1220
        preserve=[in_pkey_col], start=0))
1221
    input_joins = [in_table]
1222
    db.log_debug('Temp table: '+strings.as_tt(in_table.to_str(db)), level=2)
1223
    
1224
    mapping = sql_gen.ColDict(db, out_table, mapping)
1225
        # after applying dicts.join() because that returns a plain dict
1226
    
1227
    # Resolve default value column
1228
    if default != None:
1229
        try: default = mapping[default]
1230
        except KeyError:
1231
            db.log_debug('Default value column '
1232
                +strings.as_tt(strings.repr_no_u(default))
1233
                +' does not exist in mapping, falling back to None', level=2.1)
1234
            default = None
1235
    
1236
    pkeys_names = [in_pkey, out_pkey]
1237
    pkeys_cols = [in_pkey_col, out_pkey_col]
1238
    
1239
    pkeys_table_exists_ref = [False]
1240
    def insert_into_pkeys(joins, cols):
1241
        query = mk_select(db, joins, cols, order_by=None, start=0)
1242
        if pkeys_table_exists_ref[0]:
1243
            insert_select(db, into, pkeys_names, query)
1244
        else:
1245
            run_query_into(db, query, into=into)
1246
            pkeys_table_exists_ref[0] = True
1247
    
1248
    limit_ref = [None]
1249
    conds = set()
1250
    distinct_on = sql_gen.ColDict(db, out_table)
1251
    def mk_main_select(joins, cols):
1252
        distinct_on_cols = [c.to_Col() for c in distinct_on.values()]
1253
        return mk_select(db, joins, cols, conds, distinct_on_cols,
1254
            limit=limit_ref[0], start=0)
1255
    
1256
    exc_strs = set()
1257
    def log_exc(e):
1258
        e_str = exc.str_(e, first_line_only=True)
1259
        log_debug('Caught exception: '+e_str)
1260
        assert e_str not in exc_strs # avoid infinite loops
1261
        exc_strs.add(e_str)
1262
    
1263
    def remove_all_rows():
1264
        log_debug('Ignoring all rows')
1265
        limit_ref[0] = 0 # just create an empty pkeys table
1266
    
1267
    def ignore(in_col, value, e):
1268
        track_data_error(db, errors_table_, in_col.srcs, value, e.cause.pgcode,
1269
            e.cause.pgerror)
1270
        log_debug('Ignoring rows with '+strings.as_tt(repr(in_col))+' = '
1271
            +strings.as_tt(repr(value)))
1272
    
1273
    def remove_rows(in_col, value, e):
1274
        ignore(in_col, value, e)
1275
        cond = (in_col, sql_gen.CompareCond(value, '!='))
1276
        assert cond not in conds # avoid infinite loops
1277
        conds.add(cond)
1278
    
1279
    def invalid2null(in_col, value, e):
1280
        ignore(in_col, value, e)
1281
        update(db, in_table, [(in_col, None)],
1282
            sql_gen.ColValueCond(in_col, value))
1283
    
1284
    def insert_pkeys_table(which):
1285
        return sql_gen.Table(sql_gen.add_suffix(in_table.name,
1286
            '_insert_'+which+'_pkeys'))
1287
    insert_out_pkeys = insert_pkeys_table('out')
1288
    insert_in_pkeys = insert_pkeys_table('in')
1289
    
1290
    # Do inserts and selects
1291
    join_cols = sql_gen.ColDict(db, out_table)
1292
    while True:
1293
        if limit_ref[0] == 0: # special case
1294
            log_debug('Creating an empty pkeys table')
1295
            cur = run_query_into(db, mk_select(db, out_table, [out_pkey],
1296
                limit=limit_ref[0]), into=insert_out_pkeys)
1297
            break # don't do main case
1298
        
1299
        has_joins = join_cols != {}
1300
        
1301
        log_debug('Trying to insert new rows')
1302
        
1303
        # Prepare to insert new rows
1304
        insert_joins = input_joins[:] # don't modify original!
1305
        insert_args = dict(recover=True, cacheable=False)
1306
        if has_joins:
1307
            insert_joins.append(sql_gen.Join(out_table, join_cols,
1308
                sql_gen.filter_out))
1309
        else:
1310
            insert_args.update(dict(returning=out_pkey, into=insert_out_pkeys))
1311
        main_select = mk_main_select(insert_joins, mapping.values())
1312
        
1313
        try:
1314
            cur = insert_select(db, out_table, mapping.keys(), main_select,
1315
                **insert_args)
1316
            break # insert successful
1317
        except DuplicateKeyException, e:
1318
            log_exc(e)
1319
            
1320
            old_join_cols = join_cols.copy()
1321
            distinct_on.update(util.dict_subset(mapping, e.cols))
1322
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
1323
            log_debug('Ignoring existing rows, comparing on these columns:\n'
1324
                +strings.as_inline_table(join_cols, ustr=col_ustr))
1325
            assert join_cols != old_join_cols # avoid infinite loops
1326
        except NullValueException, e:
1327
            log_exc(e)
1328
            
1329
            out_col, = e.cols
1330
            try: in_col = mapping[out_col]
1331
            except KeyError:
1332
                log_debug('Missing mapping for NOT NULL column '+out_col)
1333
                remove_all_rows()
1334
            else: remove_rows(in_col, None, e)
1335
        except FunctionValueException, e:
1336
            log_exc(e)
1337
            
1338
            func_name = e.name
1339
            value = e.value
1340
            for out_col, in_col in mapping.iteritems():
1341
                in_col = sql_gen.unwrap_func_call(in_col, func_name)
1342
                invalid2null(in_col, value, e)
1343
        except MissingCastException, e:
1344
            log_exc(e)
1345
            
1346
            out_col = e.col
1347
            type_ = e.type
1348
            
1349
            log_debug('Casting '+strings.as_tt(out_col)+' input to '
1350
                +strings.as_tt(type_))
1351
            def wrap_func(col): return cast(db, type_, col, errors_table_)
1352
            mapping[out_col] = sql_gen.wrap(wrap_func, mapping[out_col])
1353
        except DatabaseErrors, e:
1354
            log_exc(e)
1355
            
1356
            log_debug('No handler for exception')
1357
            on_error(e)
1358
            remove_all_rows()
1359
        # after exception handled, rerun loop with additional constraints
1360
    
1361
    if row_ct_ref != None and cur.rowcount >= 0:
1362
        row_ct_ref[0] += cur.rowcount
1363
    
1364
    if has_joins:
1365
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
1366
        log_debug('Getting output table pkeys of existing/inserted rows')
1367
        insert_into_pkeys(select_joins, pkeys_cols)
1368
    else:
1369
        add_row_num(db, insert_out_pkeys) # for joining with input pkeys
1370
        
1371
        log_debug('Getting input table pkeys of inserted rows')
1372
        run_query_into(db, mk_main_select(input_joins, [in_pkey]),
1373
            into=insert_in_pkeys)
1374
        add_row_num(db, insert_in_pkeys) # for joining with output pkeys
1375
        
1376
        assert table_row_count(db, insert_out_pkeys) == table_row_count(db,
1377
            insert_in_pkeys)
1378
        
1379
        log_debug('Combining output and input pkeys in inserted order')
1380
        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,
1381
            {row_num_col: sql_gen.join_same_not_null})]
1382
        insert_into_pkeys(pkey_joins, pkeys_names)
1383
    
1384
    db.log_debug('Adding pkey on pkeys table to enable fast joins', level=2.5)
1385
    add_pkey(db, into)
1386
    
1387
    log_debug('Setting pkeys of missing rows to '+strings.as_tt(repr(default)))
1388
    missing_rows_joins = input_joins+[sql_gen.Join(into,
1389
        {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
1390
        # must use join_same_not_null or query will take forever
1391
    insert_into_pkeys(missing_rows_joins,
1392
        [in_pkey_col, sql_gen.NamedCol(out_pkey, default)])
1393
    
1394
    assert table_row_count(db, into) == table_row_count(db, in_table)
1395
    
1396
    srcs = []
1397
    if is_func: srcs = sql_gen.cols_srcs(in_cols)
1398
    return sql_gen.Col(out_pkey, into, srcs)
1399

    
1400
##### Data cleanup
1401

    
1402
def cleanup_table(db, table, cols):
1403
    table = sql_gen.as_Table(table)
1404
    cols = map(sql_gen.as_Col, cols)
1405
    
1406
    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '
1407
        +db.esc_value(r'\N')+')')
1408
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))
1409
        for v in cols]
1410
    
1411
    update(db, table, changes)
(24-24/36)