Project

General

Profile

1
# Database access
2

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

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

    
18
##### Exceptions
19

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

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

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

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

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

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

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

    
63
class NameException(DbException): pass
64

    
65
class DuplicateKeyException(ConstraintException): pass
66

    
67
class NullValueException(ConstraintException): pass
68

    
69
class FunctionValueException(ExceptionWithNameValue): pass
70

    
71
class DuplicateTableException(ExceptionWithName): pass
72

    
73
class DuplicateFunctionException(ExceptionWithName): pass
74

    
75
class EmptyRowException(DbException): pass
76

    
77
##### Warnings
78

    
79
class DbWarning(UserWarning): pass
80

    
81
##### Result retrieval
82

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

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

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

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

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

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

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

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

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

    
108
##### Input validation
109

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

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

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

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

    
128
##### Database connections
129

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

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

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

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

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

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

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

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

    
376
connect = DbConn
377

    
378
##### Querying
379

    
380
def run_raw_query(db, *args, **kw_args):
381
    '''For params, see DbConn.run_query()'''
382
    return db.run_query(*args, **kw_args)
383

    
384
def mogrify(db, query, params):
385
    module = util.root_module(db.db)
386
    if module == 'psycopg2': return db.db.cursor().mogrify(query, params)
387
    else: raise NotImplementedError("Can't mogrify query for "+module+
388
        ' database')
389

    
390
##### Recoverable querying
391

    
392
def with_savepoint(db, func): return db.with_savepoint(func)
393

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

    
455
##### Basic queries
456

    
457
def next_version(name):
458
    version = 1 # first existing name was version 0
459
    match = re.match(r'^(.*)#(\d+)$', name)
460
    if match:
461
        name, version = match.groups()
462
        version = int(version)+1
463
    return sql_gen.add_suffix(name, '#'+str(version))
464

    
465
def run_query_into(db, query, params, into=None, *args, **kw_args):
466
    '''Outputs a query to a temp table.
467
    For params, see run_query().
468
    '''
469
    if into == None: return run_query(db, query, params, *args, **kw_args)
470
    else: # place rows in temp table
471
        assert isinstance(into, sql_gen.Table)
472
        
473
        kw_args['recover'] = True
474
        kw_args.setdefault('log_ignore_excs', (DuplicateTableException,))
475
        
476
        temp = not db.autocommit # tables are permanent in autocommit mode
477
        # "temporary tables cannot specify a schema name", so remove schema
478
        if temp: into.schema = None
479
        
480
        while True:
481
            try:
482
                create_query = 'CREATE'
483
                if temp: create_query += ' TEMP'
484
                create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
485
                
486
                return run_query(db, create_query, params, *args, **kw_args)
487
                    # CREATE TABLE AS sets rowcount to # rows in query
488
            except DuplicateTableException, e:
489
                into.name = next_version(into.name)
490
                # try again with next version of name
491

    
492
order_by_pkey = object() # tells mk_select() to order by the pkey
493

    
494
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
495

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

    
575
def select(db, *args, **kw_args):
576
    '''For params, see mk_select() and run_query()'''
577
    recover = kw_args.pop('recover', None)
578
    cacheable = kw_args.pop('cacheable', True)
579
    log_level = kw_args.pop('log_level', 2)
580
    
581
    query, params = mk_select(db, *args, **kw_args)
582
    return run_query(db, query, params, recover, cacheable, log_level=log_level)
583

    
584
def mk_insert_select(db, table, cols=None, select_query=None, params=None,
585
    returning=None, embeddable=False):
586
    '''
587
    @param returning str|None An inserted column (such as pkey) to return
588
    @param embeddable Whether the query should be embeddable as a nested SELECT.
589
        Warning: If you set this and cacheable=True when the query is run, the
590
        query will be fully cached, not just if it raises an exception.
591
    '''
592
    table = sql_gen.as_Table(table)
593
    if cols == []: cols = None # no cols (all defaults) = unknown col names
594
    if cols != None:
595
        cols = [sql_gen.to_name_only_col(v, table).to_str(db) for v in cols]
596
    if select_query == None: select_query = 'DEFAULT VALUES'
597
    if returning != None: returning = sql_gen.as_Col(returning, table)
598
    
599
    # Build query
600
    first_line = 'INSERT INTO '+table.to_str(db)
601
    query = first_line
602
    if cols != None: query += '\n('+', '.join(cols)+')'
603
    query += '\n'+select_query
604
    
605
    if returning != None:
606
        returning_name = copy.copy(returning)
607
        returning_name.table = None
608
        returning_name = returning_name.to_str(db)
609
        query += '\nRETURNING '+returning_name
610
    
611
    if embeddable:
612
        assert returning != None
613
        
614
        # Create function
615
        function_name = sql_gen.clean_name(first_line)
616
        return_type = 'SETOF '+returning.to_str(db)+'%TYPE'
617
        while True:
618
            try:
619
                function = sql_gen.TempFunction(function_name, db.autocommit)
620
                
621
                function_query = '''\
622
CREATE FUNCTION '''+function.to_str(db)+'''()
623
RETURNS '''+return_type+'''
624
LANGUAGE sql
625
AS $$
626
'''+mogrify(db, query, params)+''';
627
$$;
628
'''
629
                run_query(db, function_query, recover=True, cacheable=True,
630
                    log_ignore_excs=(DuplicateFunctionException,))
631
                break # this version was successful
632
            except DuplicateFunctionException, e:
633
                function_name = next_version(function_name)
634
                # try again with next version of name
635
        
636
        # Return query that uses function
637
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(function),
638
            [returning_name]) # AS clause requires function alias
639
        return mk_select(db, func_table, start=0, order_by=None)
640
    
641
    return (query, params)
642

    
643
def insert_select(db, *args, **kw_args):
644
    '''For params, see mk_insert_select() and run_query_into()
645
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
646
        values in
647
    '''
648
    into = kw_args.pop('into', None)
649
    if into != None: kw_args['embeddable'] = True
650
    recover = kw_args.pop('recover', None)
651
    cacheable = kw_args.pop('cacheable', True)
652
    log_level = kw_args.pop('log_level', 2)
653
    
654
    query, params = mk_insert_select(db, *args, **kw_args)
655
    return run_query_into(db, query, params, into, recover=recover,
656
        cacheable=cacheable, log_level=log_level)
657

    
658
default = object() # tells insert() to use the default value for a column
659

    
660
def insert(db, table, row, *args, **kw_args):
661
    '''For params, see insert_select()'''
662
    if lists.is_seq(row): cols = None
663
    else:
664
        cols = row.keys()
665
        row = row.values()
666
    row = list(row) # ensure that "!= []" works
667
    
668
    # Check for special values
669
    labels = []
670
    values = []
671
    for value in row:
672
        value = sql_gen.remove_col_rename(sql_gen.as_Value(value)).value
673
        if value is default: labels.append('DEFAULT')
674
        else:
675
            labels.append('%s')
676
            values.append(value)
677
    
678
    # Build query
679
    if values != []: query = 'VALUES ('+(', '.join(labels))+')'
680
    else: query = None
681
    
682
    return insert_select(db, table, cols, query, values, *args, **kw_args)
683

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

    
700
def update(db, *args, **kw_args):
701
    '''For params, see mk_update() and run_query()'''
702
    recover = kw_args.pop('recover', None)
703
    
704
    return run_query(db, mk_update(db, *args, **kw_args), [], recover)
705

    
706
def last_insert_id(db):
707
    module = util.root_module(db.db)
708
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
709
    elif module == 'MySQLdb': return db.insert_id()
710
    else: return None
711

    
712
def truncate(db, table, schema='public'):
713
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
714

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

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

    
757
def cast(db, type_, col, errors_table=None):
758
    '''Casts an (unrenamed) column or value.
759
    If a column, converts any errors to warnings.
760
    @param col sql_gen.Col|sql_gen.Literal
761
    @param errors_table None|sql_gen.Table|str
762
        If set and col is a column with srcs, saves any errors in this table,
763
        using column's srcs attr as the source columns.
764
    '''
765
    if isinstance(col, sql_gen.Literal): # literal value, so just cast
766
        return sql_gen.CustomCode(col.to_str(db)+'::'+type_)
767
    
768
    assert isinstance(col, sql_gen.Col)
769
    assert not isinstance(col, sql_gen.NamedCol)
770
    
771
    save_errors = errors_table != None and col.srcs != ()
772
    if save_errors:
773
        errors_table = sql_gen.as_Table(errors_table)
774
        srcs = map(sql_gen.to_name_only_col, col.srcs)
775
        function_name = str(sql_gen.FunctionCall(type_, *srcs))
776
    else: function_name = type_
777
    function = sql_gen.TempFunction(function_name, db.autocommit)
778
    
779
    while True:
780
        # Create function definition
781
        query = '''\
782
CREATE FUNCTION '''+function.to_str(db)+'''(value text)
783
RETURNS '''+type_+'''
784
LANGUAGE plpgsql
785
'''
786
        if not save_errors: query += 'IMMUTABLE\n'
787
        query += '''\
788
STRICT
789
AS $$
790
BEGIN
791
    /* The explicit cast to the return type is needed to make the cast happen
792
    inside the try block. (Implicit casts to the return type happen at the end
793
    of the function, outside any block.) */
794
    RETURN value::'''+type_+''';
795
EXCEPTION
796
    WHEN data_exception THEN
797
'''
798
        if save_errors:
799
            cols = [sql_gen.Literal(c.name) for c in srcs]
800
            query += '''\
801
        -- Save error in errors table.
802
        BEGIN
803
            -- Insert the value and error for *each* source column.
804
            INSERT INTO '''+errors_table.to_str(db)+'''
805
            ("column", value, error_code, error)
806
            SELECT *
807
            FROM (VALUES '''+(', '.join(('('+c.to_str(db)+')' for c in cols))
808
                )+''') AS c
809
            CROSS JOIN (VALUES (value, SQLSTATE, SQLERRM)) AS v
810
            ;
811
        EXCEPTION
812
            WHEN unique_violation THEN NULL; -- ignore duplicate key
813
        END;
814
        
815
'''
816
        query += '''\
817
        RAISE WARNING '%', SQLERRM;
818
        RETURN NULL;
819
END;
820
$$;
821
'''
822
        
823
        # Create function
824
        try:
825
            run_query(db, query, recover=True, cacheable=True,
826
                log_ignore_excs=(DuplicateFunctionException,))
827
            break # successful
828
        except DuplicateFunctionException:
829
            if save_errors: function.name = next_version(function.name)
830
                # try again with next version of name
831
            else: break # plain cast function, so only need one version
832
    
833
    return sql_gen.FunctionCall(function, col)
834

    
835
##### Database structure queries
836

    
837
def table_row_count(db, table, recover=None):
838
    return value(run_query(db, *mk_select(db, table, [sql_gen.row_count],
839
        order_by=None, start=0), recover=recover, log_level=3))
840

    
841
def table_cols(db, table, recover=None):
842
    return list(col_names(select(db, table, limit=0, order_by=None,
843
        recover=recover, log_level=4)))
844

    
845
def pkey(db, table, recover=None):
846
    '''Assumed to be first column in table'''
847
    return table_cols(db, table, recover)[0]
848

    
849
not_null_col = 'not_null_col'
850

    
851
def table_not_null_col(db, table, recover=None):
852
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
853
    if not_null_col in table_cols(db, table, recover): return not_null_col
854
    else: return pkey(db, table, recover)
855

    
856
def index_cols(db, table, index):
857
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
858
    automatically created. When you don't know whether something is a UNIQUE
859
    constraint or a UNIQUE index, use this function.'''
860
    module = util.root_module(db.db)
861
    if module == 'psycopg2':
862
        return list(values(run_query(db, '''\
863
SELECT attname
864
FROM
865
(
866
        SELECT attnum, attname
867
        FROM pg_index
868
        JOIN pg_class index ON index.oid = indexrelid
869
        JOIN pg_class table_ ON table_.oid = indrelid
870
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
871
        WHERE
872
            table_.relname = %(table)s
873
            AND index.relname = %(index)s
874
    UNION
875
        SELECT attnum, attname
876
        FROM
877
        (
878
            SELECT
879
                indrelid
880
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
881
                    AS indkey
882
            FROM pg_index
883
            JOIN pg_class index ON index.oid = indexrelid
884
            JOIN pg_class table_ ON table_.oid = indrelid
885
            WHERE
886
                table_.relname = %(table)s
887
                AND index.relname = %(index)s
888
        ) s
889
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
890
) s
891
ORDER BY attnum
892
''',
893
            {'table': table, 'index': index}, cacheable=True, log_level=4)))
894
    else: raise NotImplementedError("Can't list index columns for "+module+
895
        ' database')
896

    
897
def constraint_cols(db, table, constraint):
898
    module = util.root_module(db.db)
899
    if module == 'psycopg2':
900
        return list(values(run_query(db, '''\
901
SELECT attname
902
FROM pg_constraint
903
JOIN pg_class ON pg_class.oid = conrelid
904
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
905
WHERE
906
    relname = %(table)s
907
    AND conname = %(constraint)s
908
ORDER BY attnum
909
''',
910
            {'table': table, 'constraint': constraint})))
911
    else: raise NotImplementedError("Can't list constraint columns for "+module+
912
        ' database')
913

    
914
row_num_col = '_row_num'
915

    
916
def add_index(db, exprs, table=None, unique=False):
917
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
918
    Currently, only function calls are supported as expressions.
919
    '''
920
    if not lists.is_seq(exprs): exprs = [exprs]
921
    
922
    # Parse exprs
923
    old_exprs = exprs[:]
924
    exprs = []
925
    cols = []
926
    for i, expr in enumerate(old_exprs):
927
        expr = copy.deepcopy(expr) # don't modify input!
928
        expr = sql_gen.as_Col(expr)
929
        
930
        # Extract col
931
        if isinstance(expr, sql_gen.FunctionCall):
932
            col = expr.args[0]
933
            expr = sql_gen.Expr(expr)
934
        else: col = expr
935
        
936
        # Extract table
937
        if table == None:
938
            assert sql_gen.is_table_col(col)
939
            table = col.table
940
        
941
        col.table = None
942
        
943
        exprs.append(expr)
944
        cols.append(col)
945
    
946
    table = sql_gen.as_Table(table)
947
    index = sql_gen.Table(str(sql_gen.Col(','.join(map(str, cols)), table)))
948
    
949
    str_ = 'CREATE'
950
    if unique: str_ += ' UNIQUE'
951
    str_ += ' INDEX '+index.to_str(db)+' ON '+table.to_str(db)+' ('+(
952
        ', '.join((v.to_str(db) for v in exprs)))+')'
953
    
954
    try: run_query(db, str_, recover=True, cacheable=True, log_level=3)
955
    except DuplicateTableException: pass # index already existed
956

    
957
def add_pkey(db, table, cols=None, recover=None):
958
    '''Adds a primary key.
959
    @param cols [sql_gen.Col,...] The columns in the primary key.
960
        Defaults to the first column in the table.
961
    @pre The table must not already have a primary key.
962
    '''
963
    table = sql_gen.as_Table(table)
964
    if cols == None: cols = [pkey(db, table, recover)]
965
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
966
    
967
    index = sql_gen.as_Table(sql_gen.add_suffix(table.name, '_pkey'))
968
    try:
969
        run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD CONSTRAINT '
970
            +index.to_str(db)+' PRIMARY KEY ('+(', '.join(col_strs))+')',
971
            recover=True, cacheable=True, log_level=3,
972
            log_ignore_excs=(DuplicateTableException,))
973
    except DuplicateTableException, e:
974
        index.name = next_version(index.name)
975
        # try again with next version of name
976

    
977
def add_row_num(db, table):
978
    '''Adds a row number column to a table. Its name is in row_num_col. It will
979
    be the primary key.'''
980
    table = sql_gen.as_Table(table).to_str(db)
981
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
982
        +' serial NOT NULL PRIMARY KEY', log_level=3)
983

    
984
def create_table(db, table, cols, has_pkey=True):
985
    '''Creates a table.
986
    @param cols [sql_gen.TypedCol,...] The column names and types
987
    @param has_pkey If set, the first column becomes the primary key.
988
    '''
989
    table = sql_gen.as_Table(table)
990
    
991
    if has_pkey:
992
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
993
        pkey.type += ' NOT NULL PRIMARY KEY'
994
    
995
    str_ = 'CREATE TABLE '+table.to_str(db)+' (\n'
996
    str_ += '\n, '.join(v.to_str(db) for v in cols)
997
    str_ += '\n);\n'
998
    run_query(db, str_, cacheable=True, log_level=2)
999

    
1000
def vacuum(db, table):
1001
    table = sql_gen.as_Table(table)
1002
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1003
        log_level=3))
1004

    
1005
def tables(db, schema_like='public', table_like='%'):
1006
    module = util.root_module(db.db)
1007
    params = {'schema_like': schema_like, 'table_like': table_like}
1008
    if module == 'psycopg2':
1009
        return values(run_query(db, '''\
1010
SELECT tablename
1011
FROM pg_tables
1012
WHERE
1013
    schemaname LIKE %(schema_like)s
1014
    AND tablename LIKE %(table_like)s
1015
ORDER BY tablename
1016
''',
1017
            params, cacheable=True))
1018
    elif module == 'MySQLdb':
1019
        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,
1020
            cacheable=True))
1021
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1022

    
1023
def table_exists(db, table):
1024
    table = sql_gen.as_Table(table)
1025
    return list(tables(db, table.schema, table.name)) != []
1026

    
1027
def errors_table(db, table, if_exists=True):
1028
    '''
1029
    @param if_exists If set, returns None if the errors table doesn't exist
1030
    @return None|sql_gen.Table
1031
    '''
1032
    table = sql_gen.as_Table(table)
1033
    if table.srcs != (): table = table.srcs[0]
1034
    
1035
    errors_table = sql_gen.suffixed_table(table, '.errors')
1036
    if if_exists and not table_exists(db, errors_table): return None
1037
    return errors_table
1038

    
1039
##### Database management
1040

    
1041
def empty_db(db, schema='public', **kw_args):
1042
    '''For kw_args, see tables()'''
1043
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1044

    
1045
##### Heuristic queries
1046

    
1047
def put(db, table, row, pkey_=None, row_ct_ref=None):
1048
    '''Recovers from errors.
1049
    Only works under PostgreSQL (uses INSERT RETURNING).
1050
    '''
1051
    row = sql_gen.ColDict(db, table, row)
1052
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
1053
    
1054
    try:
1055
        cur = insert(db, table, row, pkey_, recover=True)
1056
        if row_ct_ref != None and cur.rowcount >= 0:
1057
            row_ct_ref[0] += cur.rowcount
1058
        return value(cur)
1059
    except DuplicateKeyException, e:
1060
        row = sql_gen.ColDict(db, table,
1061
            util.dict_subset_right_join(row, e.cols))
1062
        return value(select(db, table, [pkey_], row, recover=True))
1063

    
1064
def get(db, table, row, pkey, row_ct_ref=None, create=False):
1065
    '''Recovers from errors'''
1066
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
1067
    except StopIteration:
1068
        if not create: raise
1069
        return put(db, table, row, pkey, row_ct_ref) # insert new row
1070

    
1071
def is_func_result(col):
1072
    return col.table.name.find('(') >= 0 and col.name == 'result'
1073

    
1074
def into_table_name(out_table, in_tables0, mapping, is_func):
1075
    def in_col_str(in_col):
1076
        in_col = sql_gen.remove_col_rename(in_col)
1077
        if isinstance(in_col, sql_gen.Col):
1078
            table = in_col.table
1079
            if table == in_tables0:
1080
                in_col = sql_gen.to_name_only_col(in_col)
1081
            elif is_func_result(in_col): in_col = table # omit col name
1082
        return str(in_col)
1083
    
1084
    str_ = str(out_table)
1085
    if is_func:
1086
        str_ += '('
1087
        
1088
        try: value_in_col = mapping['value']
1089
        except KeyError:
1090
            str_ += ', '.join((str(k)+'='+in_col_str(v)
1091
                for k, v in mapping.iteritems()))
1092
        else: str_ += in_col_str(value_in_col)
1093
        
1094
        str_ += ')'
1095
    else: str_ += '_pkeys'
1096
    return str_
1097

    
1098
def put_table(db, out_table, in_tables, mapping, row_ct_ref=None, into=None,
1099
    default=None, is_func=False):
1100
    '''Recovers from errors.
1101
    Only works under PostgreSQL (uses INSERT RETURNING).
1102
    @param in_tables The main input table to select from, followed by a list of
1103
        tables to join with it using the main input table's pkey
1104
    @param mapping dict(out_table_col=in_table_col, ...)
1105
        * out_table_col: str (*not* sql_gen.Col)
1106
        * in_table_col: sql_gen.Col|literal-value
1107
    @param into The table to contain the output and input pkeys.
1108
        Defaults to `out_table.name+'_pkeys'`.
1109
    @param default The *output* column to use as the pkey for missing rows.
1110
        If this output column does not exist in the mapping, uses None.
1111
    @param is_func Whether out_table is the name of a SQL function, not a table
1112
    @return sql_gen.Col Where the output pkeys are made available
1113
    '''
1114
    out_table = sql_gen.as_Table(out_table)
1115
    
1116
    def log_debug(msg): db.log_debug(msg, level=1.5)
1117
    def col_ustr(str_):
1118
        return strings.repr_no_u(sql_gen.remove_col_rename(str_))
1119
    
1120
    log_debug('********** New iteration **********')
1121
    log_debug('Inserting these input columns into '+strings.as_tt(
1122
        out_table.to_str(db))+':\n'+strings.as_table(mapping, ustr=col_ustr))
1123
    
1124
    # Create input joins from list of input tables
1125
    in_tables_ = in_tables[:] # don't modify input!
1126
    in_tables0 = in_tables_.pop(0) # first table is separate
1127
    in_pkey = pkey(db, in_tables0, recover=True)
1128
    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
1129
    input_joins = [in_tables0]+[sql_gen.Join(v,
1130
        {in_pkey: sql_gen.join_same_not_null}) for v in in_tables_]
1131
    
1132
    if into == None:
1133
        into = into_table_name(out_table, in_tables0, mapping, is_func)
1134
    into = sql_gen.as_Table(into)
1135
    
1136
    # Set column sources
1137
    in_cols = filter(sql_gen.is_table_col, mapping.values())
1138
    for col in in_cols:
1139
        if col.table == in_tables0: col.set_srcs(sql_gen.src_self)
1140
    
1141
    log_debug('Joining together input tables into temp table')
1142
    # Place in new table for speed and so don't modify input if values edited
1143
    in_table = sql_gen.Table('in')
1144
    mapping = dicts.join(mapping, flatten(db, in_table, input_joins, in_cols,
1145
        preserve=[in_pkey_col], start=0))
1146
    input_joins = [in_table]
1147
    db.log_debug('Temp table: '+strings.as_tt(in_table.to_str(db)), level=2)
1148
    
1149
    mapping = sql_gen.ColDict(db, out_table, mapping)
1150
        # after applying dicts.join() because that returns a plain dict
1151
    
1152
    # Resolve default value column
1153
    try: default = mapping[default]
1154
    except KeyError:
1155
        if default != None:
1156
            db.log_debug('Default value column '
1157
                +strings.as_tt(strings.repr_no_u(default))
1158
                +' does not exist in mapping, falling back to None', level=2.1)
1159
            default = None
1160
    
1161
    out_pkey = pkey(db, out_table, recover=True)
1162
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
1163
    
1164
    pkeys_names = [in_pkey, out_pkey]
1165
    pkeys_cols = [in_pkey_col, out_pkey_col]
1166
    
1167
    pkeys_table_exists_ref = [False]
1168
    def insert_into_pkeys(joins, cols):
1169
        query, params = mk_select(db, joins, cols, order_by=None, start=0)
1170
        if pkeys_table_exists_ref[0]:
1171
            insert_select(db, into, pkeys_names, query, params)
1172
        else:
1173
            run_query_into(db, query, params, into=into)
1174
            pkeys_table_exists_ref[0] = True
1175
    
1176
    limit_ref = [None]
1177
    conds = set()
1178
    distinct_on = []
1179
    def mk_main_select(joins, cols):
1180
        return mk_select(db, joins, cols, conds, distinct_on,
1181
            limit=limit_ref[0], start=0)
1182
    
1183
    exc_strs = set()
1184
    def log_exc(e):
1185
        e_str = exc.str_(e, first_line_only=True)
1186
        log_debug('Caught exception: '+e_str)
1187
        assert e_str not in exc_strs # avoid infinite loops
1188
        exc_strs.add(e_str)
1189
    def remove_all_rows():
1190
        log_debug('Returning NULL for all rows')
1191
        limit_ref[0] = 0 # just create an empty pkeys table
1192
    def ignore(in_col, value):
1193
        in_col_str = strings.as_tt(repr(in_col))
1194
        db.log_debug('Adding index on '+in_col_str+' to enable fast filtering',
1195
            level=2.5)
1196
        add_index(db, in_col)
1197
        log_debug('Ignoring rows with '+in_col_str+' = '
1198
            +strings.as_tt(repr(value)))
1199
    def remove_rows(in_col, value):
1200
        ignore(in_col, value)
1201
        cond = (in_col, sql_gen.CompareCond(value, '!='))
1202
        assert cond not in conds # avoid infinite loops
1203
        conds.add(cond)
1204
    def invalid2null(in_col, value):
1205
        ignore(in_col, value)
1206
        update(db, in_table, [(in_col, None)],
1207
            sql_gen.ColValueCond(in_col, value))
1208
    
1209
    def insert_pkeys_table(which):
1210
        return sql_gen.Table(sql_gen.add_suffix(in_table.name,
1211
            '_insert_'+which+'_pkeys'))
1212
    insert_out_pkeys = insert_pkeys_table('out')
1213
    insert_in_pkeys = insert_pkeys_table('in')
1214
    
1215
    # Do inserts and selects
1216
    join_cols = sql_gen.ColDict(db, out_table)
1217
    while True:
1218
        if limit_ref[0] == 0: # special case
1219
            log_debug('Creating an empty pkeys table')
1220
            cur = run_query_into(db, *mk_select(db, out_table, [out_pkey],
1221
                limit=limit_ref[0]), into=insert_out_pkeys)
1222
            break # don't do main case
1223
        
1224
        has_joins = join_cols != {}
1225
        
1226
        # Prepare to insert new rows
1227
        insert_joins = input_joins[:] # don't modify original!
1228
        insert_args = dict(recover=True, cacheable=False)
1229
        if has_joins:
1230
            distinct_on = [v.to_Col() for v in join_cols.values()]
1231
            insert_joins.append(sql_gen.Join(out_table, join_cols,
1232
                sql_gen.filter_out))
1233
        else:
1234
            insert_args.update(dict(returning=out_pkey, into=insert_out_pkeys))
1235
        main_select = mk_main_select(insert_joins, mapping.values())[0]
1236
        
1237
        log_debug('Trying to insert new rows')
1238
        try:
1239
            cur = insert_select(db, out_table, mapping.keys(), main_select,
1240
                **insert_args)
1241
            break # insert successful
1242
        except DuplicateKeyException, e:
1243
            log_exc(e)
1244
            
1245
            old_join_cols = join_cols.copy()
1246
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
1247
            log_debug('Ignoring existing rows, comparing on these columns:\n'
1248
                +strings.as_inline_table(join_cols, ustr=col_ustr))
1249
            assert join_cols != old_join_cols # avoid infinite loops
1250
        except NullValueException, e:
1251
            log_exc(e)
1252
            
1253
            out_col, = e.cols
1254
            try: in_col = mapping[out_col]
1255
            except KeyError:
1256
                log_debug('Missing mapping for NOT NULL column '+out_col)
1257
                remove_all_rows()
1258
            else: remove_rows(in_col, None)
1259
        except FunctionValueException, e:
1260
            log_exc(e)
1261
            
1262
            func_name = e.name
1263
            value = e.value
1264
            for out_col, in_col in mapping.iteritems():
1265
                invalid2null(sql_gen.unwrap_func_call(in_col, func_name), value)
1266
        except MissingCastException, e:
1267
            log_exc(e)
1268
            
1269
            out_col = e.col
1270
            type_ = e.type
1271
            
1272
            log_debug('Casting '+strings.as_tt(out_col)+' input to '
1273
                +strings.as_tt(type_))
1274
            errors_table = sql_gen.suffixed_table(in_tables0.srcs[0], '.errors')
1275
            def wrap_func(col): return cast(db, type_, col, errors_table)
1276
            mapping[out_col] = sql_gen.wrap(wrap_func, mapping[out_col])
1277
        except DatabaseErrors, e:
1278
            log_exc(e)
1279
            
1280
            msg = 'No handler for exception: '+exc.str_(e)
1281
            warnings.warn(DbWarning(msg))
1282
            log_debug(msg)
1283
            remove_all_rows()
1284
        # after exception handled, rerun loop with additional constraints
1285
    
1286
    if row_ct_ref != None and cur.rowcount >= 0:
1287
        row_ct_ref[0] += cur.rowcount
1288
    
1289
    if has_joins:
1290
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
1291
        log_debug('Getting output table pkeys of existing/inserted rows')
1292
        insert_into_pkeys(select_joins, pkeys_cols)
1293
    else:
1294
        add_row_num(db, insert_out_pkeys) # for joining with input pkeys
1295
        
1296
        log_debug('Getting input table pkeys of inserted rows')
1297
        run_query_into(db, *mk_main_select(input_joins, [in_pkey]),
1298
            into=insert_in_pkeys)
1299
        add_row_num(db, insert_in_pkeys) # for joining with output pkeys
1300
        
1301
        assert table_row_count(db, insert_out_pkeys) == table_row_count(db,
1302
            insert_in_pkeys)
1303
        
1304
        log_debug('Combining output and input pkeys in inserted order')
1305
        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,
1306
            {row_num_col: sql_gen.join_same_not_null})]
1307
        insert_into_pkeys(pkey_joins, pkeys_names)
1308
    
1309
    db.log_debug('Adding pkey on pkeys table to enable fast joins', level=2.5)
1310
    add_pkey(db, into)
1311
    
1312
    log_debug('Setting pkeys of missing rows to '+strings.as_tt(repr(default)))
1313
    missing_rows_joins = input_joins+[sql_gen.Join(into,
1314
        {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
1315
        # must use join_same_not_null or query will take forever
1316
    insert_into_pkeys(missing_rows_joins,
1317
        [in_pkey_col, sql_gen.NamedCol(out_pkey, default)])
1318
    
1319
    assert table_row_count(db, into) == table_row_count(db, in_table)
1320
    
1321
    srcs = []
1322
    if is_func: srcs = sql_gen.cols_srcs(in_cols)
1323
    return sql_gen.Col(out_pkey, into, srcs)
1324

    
1325
##### Data cleanup
1326

    
1327
def cleanup_table(db, table, cols):
1328
    table = sql_gen.as_Table(table)
1329
    cols = map(sql_gen.as_Col, cols)
1330
    
1331
    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '
1332
        +db.esc_value(r'\N')+')')
1333
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))
1334
        for v in cols]
1335
    
1336
    update(db, table, changes)
(24-24/36)