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: '+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
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
            if params == [] or params == (): params = None
216
                # None turns off the mogrifier
217
            
218
            self._is_insert = query.startswith('INSERT')
219
            self.query_lookup = _query_lookup(query, params)
220
            try:
221
                try:
222
                    cur = self.inner.execute(query, params)
223
                    self.outer.do_autocommit()
224
                finally: self.query = get_cur_query(self.inner, query, params)
225
            except Exception, e:
226
                _add_cursor_info(e, self, query, params)
227
                self.result = e # cache the exception as the result
228
                self._cache_result()
229
                raise
230
            # Fetch all rows so result will be cached
231
            if self.rowcount == 0 and not self._is_insert: consume_rows(self)
232
            return cur
233
        
234
        def fetchone(self):
235
            row = self.inner.fetchone()
236
            if row != None: self.result.append(row)
237
            # otherwise, fetched all rows
238
            else: self._cache_result()
239
            return row
240
        
241
        def _cache_result(self):
242
            # For inserts, only cache exceptions since inserts are not
243
            # idempotent, but an invalid insert will always be invalid
244
            if self.query_results != None and (not self._is_insert
245
                or isinstance(self.result, Exception)):
246
                
247
                assert self.query_lookup != None
248
                self.query_results[self.query_lookup] = self.CacheCursor(
249
                    util.dict_subset(dicts.AttrsDictView(self),
250
                    ['query', 'result', 'rowcount', 'description']))
251
        
252
        class CacheCursor:
253
            def __init__(self, cached_result): self.__dict__ = cached_result
254
            
255
            def execute(self, *args, **kw_args):
256
                if isinstance(self.result, Exception): raise self.result
257
                # otherwise, result is a rows list
258
                self.iter = iter(self.result)
259
            
260
            def fetchone(self):
261
                try: return self.iter.next()
262
                except StopIteration: return None
263
    
264
    def esc_value(self, value):
265
        try: str_ = self.mogrify('%s', [value])
266
        except NotImplementedError, e:
267
            module = util.root_module(self.db)
268
            if module == 'MySQLdb':
269
                import _mysql
270
                str_ = _mysql.escape_string(value)
271
            else: raise e
272
        return strings.to_unicode(str_)
273
    
274
    def esc_name(self, name): return esc_name(self, name) # calls global func
275
    
276
    def can_mogrify(self):
277
        module = util.root_module(self.db)
278
        return module == 'psycopg2'
279
    
280
    def mogrify(self, query, params=None):
281
        if self.can_mogrify(): return self.db.cursor().mogrify(query, params)
282
        else: raise NotImplementedError("Can't mogrify query")
283
    
284
    def print_notices(self):
285
        if hasattr(self.db, 'notices'):
286
            for msg in self.db.notices:
287
                if msg not in self._notices_seen:
288
                    self._notices_seen.add(msg)
289
                    self.log_debug(msg, level=2)
290
    
291
    def run_query(self, query, params=None, cacheable=False, log_level=2,
292
        debug_msg_ref=None):
293
        '''
294
        @param log_ignore_excs The log_level will be increased by 2 if the query
295
            throws one of these exceptions.
296
        @param debug_msg_ref If specified, the log message will be returned in
297
            this instead of being output. This allows you to filter log messages
298
            depending on the result of the query.
299
        '''
300
        assert query != None
301
        
302
        if not self.caching: cacheable = False
303
        used_cache = False
304
        
305
        def log_msg(query):
306
            if used_cache: cache_status = 'cache hit'
307
            elif cacheable: cache_status = 'cache miss'
308
            else: cache_status = 'non-cacheable'
309
            return 'DB query: '+cache_status+':\n'+strings.as_code(query, 'SQL')
310
        
311
        try:
312
            # Get cursor
313
            if cacheable:
314
                query_lookup = _query_lookup(query, params)
315
                try:
316
                    cur = self.query_results[query_lookup]
317
                    used_cache = True
318
                except KeyError: cur = self.DbCursor(self)
319
            else: cur = self.db.cursor()
320
            
321
            # Log query
322
            if self.debug and debug_msg_ref == None: # log before running
323
                self.log_debug(log_msg(query), log_level)
324
            
325
            # Run query
326
            cur.execute(query, params)
327
        finally:
328
            self.print_notices()
329
            if self.debug and debug_msg_ref != None: # return after running
330
                debug_msg_ref[0] = log_msg(str(get_cur_query(cur, query,
331
                    params)))
332
        
333
        return cur
334
    
335
    def is_cached(self, query, params=None):
336
        return _query_lookup(query, params) in self.query_results
337
    
338
    def with_autocommit(self, func, autocommit=True):
339
        prev_autocommit = self.db.autocommit
340
        self.db.autocommit = autocommit
341
        try: return func()
342
        finally: self.db.autocommit = prev_autocommit
343
    
344
    def with_savepoint(self, func):
345
        savepoint = 'level_'+str(self._savepoint)
346
        self.run_query('SAVEPOINT '+savepoint, log_level=4)
347
        self._savepoint += 1
348
        try:
349
            try: return_val = func()
350
            finally:
351
                self._savepoint -= 1
352
                assert self._savepoint >= 0
353
        except:
354
            self.run_query('ROLLBACK TO SAVEPOINT '+savepoint, log_level=4)
355
            raise
356
        else:
357
            self.run_query('RELEASE SAVEPOINT '+savepoint, log_level=4)
358
            self.do_autocommit()
359
            return return_val
360
    
361
    def do_autocommit(self):
362
        '''Autocommits if outside savepoint'''
363
        assert self._savepoint >= 0
364
        if self.autocommit and self._savepoint == 0:
365
            self.log_debug('Autocommitting')
366
            self.db.commit()
367
    
368
    def col_default(self, col):
369
        table = sql_gen.Table('columns', 'information_schema')
370
        
371
        conds = [('table_name', col.table.name), ('column_name', col.name)]
372
        schema = col.table.schema
373
        if schema != None: conds.append(('table_schema', schema))
374
        
375
        return sql_gen.as_Code(value(select(self, table, ['column_default'],
376
            conds, order_by='table_schema', limit=1, log_level=3)))
377
            # TODO: order_by search_path schema order
378

    
379
connect = DbConn
380

    
381
##### Querying
382

    
383
def run_raw_query(db, *args, **kw_args):
384
    '''For params, see DbConn.run_query()'''
385
    return db.run_query(*args, **kw_args)
386

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

    
393
##### Recoverable querying
394

    
395
def with_savepoint(db, func): return db.with_savepoint(func)
396

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

    
458
##### Basic queries
459

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

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

    
501
order_by_pkey = object() # tells mk_select() to order by the pkey
502

    
503
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
504

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

    
586
def select(db, *args, **kw_args):
587
    '''For params, see mk_select() and run_query()'''
588
    recover = kw_args.pop('recover', None)
589
    cacheable = kw_args.pop('cacheable', True)
590
    log_level = kw_args.pop('log_level', 2)
591
    
592
    query, params = mk_select(db, *args, **kw_args)
593
    return run_query(db, query, params, recover, cacheable, log_level=log_level)
594

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

    
651
def insert_select(db, *args, **kw_args):
652
    '''For params, see mk_insert_select() and run_query_into()
653
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
654
        values in
655
    '''
656
    into = kw_args.pop('into', None)
657
    if into != None: kw_args['embeddable'] = True
658
    recover = kw_args.pop('recover', None)
659
    cacheable = kw_args.pop('cacheable', True)
660
    log_level = kw_args.pop('log_level', 2)
661
    
662
    query, params = mk_insert_select(db, *args, **kw_args)
663
    return run_query_into(db, query, params, into, recover=recover,
664
        cacheable=cacheable, log_level=log_level)
665

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

    
668
def insert(db, table, row, *args, **kw_args):
669
    '''For params, see insert_select()'''
670
    if lists.is_seq(row): cols = None
671
    else:
672
        cols = row.keys()
673
        row = row.values()
674
    row = list(row) # ensure that "== []" works
675
    
676
    if row == []: query = None
677
    else: query = sql_gen.Values(row).to_str(db)
678
    
679
    return insert_select(db, table, cols, query, [], *args, **kw_args)
680

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

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

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

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

    
741
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
742
    '''For params, see mk_flatten_mapping()
743
    @return See return value of mk_flatten_mapping()
744
    '''
745
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
746
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
747
    run_query_into(db, *mk_select(db, joins, cols, limit=limit, start=start),
748
        into=into)
749
    return dict(items)
750

    
751
def mk_track_data_error(db, errors_table, cols, value, error_code, error):
752
    cols = map(sql_gen.to_name_only_col, cols)
753
    
754
    columns_cols = ['column']
755
    columns = sql_gen.NamedValues('columns', columns_cols,
756
        [[c.name] for c in cols])
757
    values_cols = ['value', 'error_code', 'error']
758
    values = sql_gen.NamedValues('values', values_cols,
759
        [value, error_code, error])
760
    
761
    select_cols = columns_cols+values_cols
762
    name_only_cols = map(sql_gen.to_name_only_col, select_cols)
763
    errors_table = sql_gen.NamedTable('errors', errors_table)
764
    joins = [columns, sql_gen.Join(values, type_='CROSS'),
765
        sql_gen.Join(errors_table, dict(zip(name_only_cols, select_cols)),
766
        sql_gen.filter_out)]
767
    
768
    return mk_insert_select(db, errors_table, name_only_cols,
769
        *mk_select(db, joins, select_cols, order_by=None))[0]
770

    
771
def track_data_error(db, errors_table, *args, **kw_args):
772
    '''
773
    @param errors_table If None, does nothing.
774
    '''
775
    if errors_table == None: return
776
    run_query(db, mk_track_data_error(db, errors_table, *args, **kw_args),
777
        cacheable=True, log_level=4)
778

    
779
def cast(db, type_, col, errors_table=None):
780
    '''Casts an (unrenamed) column or value.
781
    If errors_table set and col has srcs, saves errors in errors_table (using
782
    col's srcs attr as the source columns) and converts errors to warnings.
783
    @param col sql_gen.Col|sql_gen.Literal
784
    @param errors_table None|sql_gen.Table|str
785
    '''
786
    save_errors = (errors_table != None and isinstance(col, sql_gen.Col)
787
        and col.srcs != ())
788
    if not save_errors: # can't save errors
789
        return sql_gen.CustomCode(col.to_str(db)+'::'+type_) # just cast
790
    
791
    assert not isinstance(col, sql_gen.NamedCol)
792
    
793
    errors_table = sql_gen.as_Table(errors_table)
794
    srcs = map(sql_gen.to_name_only_col, col.srcs)
795
    function_name = str(sql_gen.FunctionCall(type_, *srcs))
796
    function = sql_gen.TempFunction(function_name, db.autocommit)
797
    
798
    while True:
799
        # Create function definition
800
        query = '''\
801
CREATE FUNCTION '''+function.to_str(db)+'''(value text)
802
RETURNS '''+type_+'''
803
LANGUAGE plpgsql
804
STRICT
805
AS $$
806
BEGIN
807
    /* The explicit cast to the return type is needed to make the cast happen
808
    inside the try block. (Implicit casts to the return type happen at the end
809
    of the function, outside any block.) */
810
    RETURN value::'''+type_+''';
811
EXCEPTION
812
    WHEN data_exception THEN
813
        -- Save error in errors table.
814
        -- Insert the value and error for *each* source column.
815
'''+mk_track_data_error(db, errors_table, srcs,
816
    *map(sql_gen.CustomCode, ['value', 'SQLSTATE', 'SQLERRM']))+''';
817
        
818
        RAISE WARNING '%', SQLERRM;
819
        RETURN NULL;
820
END;
821
$$;
822
'''
823
        
824
        # Create function
825
        try:
826
            run_query(db, query, recover=True, cacheable=True,
827
                log_ignore_excs=(DuplicateFunctionException,))
828
            break # successful
829
        except DuplicateFunctionException:
830
            function.name = next_version(function.name)
831
            # try again with next version of name
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 = '''+db.esc_value(table)+'''
873
            AND index.relname = '''+db.esc_value(index)+'''
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 = '''+db.esc_value(table)+'''
887
                AND index.relname = '''+db.esc_value(index)+'''
888
        ) s
889
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
890
) s
891
ORDER BY attnum
892
'''
893
            , 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 = '''+db.esc_value(table)+'''
907
    AND conname = '''+db.esc_value(constraint)+'''
908
ORDER BY attnum
909
'''
910
            )))
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
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
968
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
969
        log_ignore_excs=(DuplicateTableException,))
970

    
971
already_indexed = object() # tells add_indexes() the pkey has already been added
972

    
973
def add_indexes(db, table, has_pkey=True):
974
    '''Adds an index on all columns in a table.
975
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
976
        index should be added on the first column.
977
        * If already_indexed, the pkey is assumed to have already been added
978
    '''
979
    cols = table_cols(db, table)
980
    if has_pkey:
981
        if has_pkey is not already_indexed: add_pkey(db, table)
982
        cols = cols[1:]
983
    for col in cols: add_index(db, col, table)
984

    
985
def add_row_num(db, table):
986
    '''Adds a row number column to a table. Its name is in row_num_col. It will
987
    be the primary key.'''
988
    table = sql_gen.as_Table(table).to_str(db)
989
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
990
        +' serial NOT NULL PRIMARY KEY', log_level=3)
991

    
992
def create_table(db, table, cols, has_pkey=True, col_indexes=True):
993
    '''Creates a table.
994
    @param cols [sql_gen.TypedCol,...] The column names and types
995
    @param has_pkey If set, the first column becomes the primary key.
996
    @param col_indexes bool|[ref]
997
        * If True, indexes will be added on all non-pkey columns.
998
        * If a list reference, [0] will be set to a function to do this.
999
          This can be used to delay index creation until the table is populated.
1000
    '''
1001
    table = sql_gen.as_Table(table)
1002
    
1003
    if has_pkey:
1004
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1005
        pkey.type += ' NOT NULL PRIMARY KEY'
1006
    
1007
    str_ = 'CREATE TABLE '+table.to_str(db)+' (\n'
1008
    str_ += '\n, '.join(v.to_str(db) for v in cols)
1009
    str_ += '\n);\n'
1010
    run_query(db, str_, cacheable=True, log_level=2)
1011
    
1012
    # Add indexes
1013
    if has_pkey: has_pkey = already_indexed
1014
    def add_indexes_(): add_indexes(db, table, has_pkey)
1015
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1016
    elif col_indexes: add_indexes_() # add now
1017

    
1018
def vacuum(db, table):
1019
    table = sql_gen.as_Table(table)
1020
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1021
        log_level=3))
1022

    
1023
def truncate(db, table, schema='public'):
1024
    table = sql_gen.as_Table(table, schema)
1025
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE')
1026

    
1027
def tables(db, schema_like='public', table_like='%', exact=False):
1028
    if exact: compare = '='
1029
    else: compare = 'LIKE'
1030
    
1031
    module = util.root_module(db.db)
1032
    params = {'schema_like': schema_like, 'table_like': table_like}
1033
    if module == 'psycopg2':
1034
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1035
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1036
        return values(select(db, 'pg_tables', ['tablename'], conds,
1037
            order_by='tablename', log_level=4))
1038
    elif module == 'MySQLdb':
1039
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1040
            , cacheable=True, log_level=4))
1041
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1042

    
1043
def table_exists(db, table):
1044
    table = sql_gen.as_Table(table)
1045
    return list(tables(db, table.schema, table.name, exact=True)) != []
1046

    
1047
def errors_table(db, table, if_exists=True):
1048
    '''
1049
    @param if_exists If set, returns None if the errors table doesn't exist
1050
    @return None|sql_gen.Table
1051
    '''
1052
    table = sql_gen.as_Table(table)
1053
    if table.srcs != (): table = table.srcs[0]
1054
    
1055
    errors_table = sql_gen.suffixed_table(table, '.errors')
1056
    if if_exists and not table_exists(db, errors_table): return None
1057
    return errors_table
1058

    
1059
##### Database management
1060

    
1061
def empty_db(db, schema='public', **kw_args):
1062
    '''For kw_args, see tables()'''
1063
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1064

    
1065
##### Heuristic queries
1066

    
1067
def put(db, table, row, pkey_=None, row_ct_ref=None):
1068
    '''Recovers from errors.
1069
    Only works under PostgreSQL (uses INSERT RETURNING).
1070
    '''
1071
    row = sql_gen.ColDict(db, table, row)
1072
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
1073
    
1074
    try:
1075
        cur = insert(db, table, row, pkey_, recover=True)
1076
        if row_ct_ref != None and cur.rowcount >= 0:
1077
            row_ct_ref[0] += cur.rowcount
1078
        return value(cur)
1079
    except DuplicateKeyException, e:
1080
        row = sql_gen.ColDict(db, table,
1081
            util.dict_subset_right_join(row, e.cols))
1082
        return value(select(db, table, [pkey_], row, recover=True))
1083

    
1084
def get(db, table, row, pkey, row_ct_ref=None, create=False):
1085
    '''Recovers from errors'''
1086
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
1087
    except StopIteration:
1088
        if not create: raise
1089
        return put(db, table, row, pkey, row_ct_ref) # insert new row
1090

    
1091
def is_func_result(col):
1092
    return col.table.name.find('(') >= 0 and col.name == 'result'
1093

    
1094
def into_table_name(out_table, in_tables0, mapping, is_func):
1095
    def in_col_str(in_col):
1096
        in_col = sql_gen.remove_col_rename(in_col)
1097
        if isinstance(in_col, sql_gen.Col):
1098
            table = in_col.table
1099
            if table == in_tables0:
1100
                in_col = sql_gen.to_name_only_col(in_col)
1101
            elif is_func_result(in_col): in_col = table # omit col name
1102
        return str(in_col)
1103
    
1104
    str_ = str(out_table)
1105
    if is_func:
1106
        str_ += '('
1107
        
1108
        try: value_in_col = mapping['value']
1109
        except KeyError:
1110
            str_ += ', '.join((str(k)+'='+in_col_str(v)
1111
                for k, v in mapping.iteritems()))
1112
        else: str_ += in_col_str(value_in_col)
1113
        
1114
        str_ += ')'
1115
    else: str_ += '_pkeys'
1116
    return str_
1117

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

    
1355
##### Data cleanup
1356

    
1357
def cleanup_table(db, table, cols):
1358
    table = sql_gen.as_Table(table)
1359
    cols = map(sql_gen.as_Col, cols)
1360
    
1361
    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '
1362
        +db.esc_value(r'\N')+')')
1363
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))
1364
        for v in cols]
1365
    
1366
    update(db, table, changes)
(24-24/36)