Project

General

Profile

1
# Database access
2

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

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

    
18
##### Exceptions
19

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

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

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

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

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

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

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

    
63
class NameException(DbException): pass
64

    
65
class DuplicateKeyException(ConstraintException): pass
66

    
67
class NullValueException(ConstraintException): pass
68

    
69
class FunctionValueException(ExceptionWithNameValue): pass
70

    
71
class DuplicateTableException(ExceptionWithName): pass
72

    
73
class DuplicateFunctionException(ExceptionWithName): pass
74

    
75
class EmptyRowException(DbException): pass
76

    
77
##### Warnings
78

    
79
class DbWarning(UserWarning): pass
80

    
81
##### Result retrieval
82

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

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

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

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

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

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

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

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

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

    
108
##### Escaping
109

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

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

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

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

    
128
##### Database connections
129

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

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

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

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

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

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

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

    
389
connect = DbConn
390

    
391
##### Recoverable querying
392

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

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

    
456
##### Basic queries
457

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

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

    
499
order_by_pkey = object() # tells mk_select() to order by the pkey
500

    
501
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
502

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

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

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

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

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

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

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

    
694
def update(db, *args, **kw_args):
695
    '''For params, see mk_update() and run_query()'''
696
    recover = kw_args.pop('recover', None)
697
    
698
    return run_query(db, mk_update(db, *args, **kw_args), recover)
699

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

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

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

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

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

    
778
def cast(db, type_, col, errors_table=None):
779
    '''Casts an (unrenamed) column or value.
780
    If errors_table set and col has srcs, saves errors in errors_table (using
781
    col's srcs attr as the source columns) and converts errors to warnings.
782
    @param col str|sql_gen.Col|sql_gen.Literal
783
    @param errors_table None|sql_gen.Table|str
784
    '''
785
    col = sql_gen.as_Col(col)
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, ensure_not_null=True):
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
    @param ensure_not_null If set, translates NULL values to sentinel values.
920
        This allows indexes to be used for comparisons where NULLs are equal.
921
    '''
922
    if not lists.is_seq(exprs): exprs = [exprs]
923
    
924
    # Parse exprs
925
    old_exprs = exprs[:]
926
    exprs = []
927
    cols = []
928
    for i, expr in enumerate(old_exprs):
929
        expr = copy.deepcopy(expr) # don't modify input!
930
        expr = sql_gen.as_Col(expr, table)
931
        
932
        # Handle nullable columns
933
        if ensure_not_null:
934
            try: expr = sql_gen.ensure_not_null(db, expr)
935
            except KeyError: pass # unknown type, so just create plain index
936
        
937
        # Extract col
938
        if isinstance(expr, sql_gen.FunctionCall):
939
            col = expr.args[0]
940
            expr = sql_gen.Expr(expr)
941
        else: col = expr
942
        assert isinstance(col, sql_gen.Col)
943
        
944
        # Extract table
945
        if table == None:
946
            assert sql_gen.is_table_col(col)
947
            table = col.table
948
        
949
        col.table = None
950
        
951
        exprs.append(expr)
952
        cols.append(col)
953
    
954
    table = sql_gen.as_Table(table)
955
    index = sql_gen.Table(str(sql_gen.Col(','.join(map(str, cols)), table)))
956
    
957
    str_ = 'CREATE'
958
    if unique: str_ += ' UNIQUE'
959
    str_ += ' INDEX '+index.to_str(db)+' ON '+table.to_str(db)+' ('+(
960
        ', '.join((v.to_str(db) for v in exprs)))+')'
961
    
962
    try: run_query(db, str_, recover=True, cacheable=True, log_level=3)
963
    except DuplicateTableException: pass # index already existed
964

    
965
def add_pkey(db, table, cols=None, recover=None):
966
    '''Adds a primary key.
967
    @param cols [sql_gen.Col,...] The columns in the primary key.
968
        Defaults to the first column in the table.
969
    @pre The table must not already have a primary key.
970
    '''
971
    table = sql_gen.as_Table(table)
972
    if cols == None: cols = [pkey(db, table, recover)]
973
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
974
    
975
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
976
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
977
        log_ignore_excs=(DuplicateTableException,))
978

    
979
already_indexed = object() # tells add_indexes() the pkey has already been added
980

    
981
def add_indexes(db, table, has_pkey=True):
982
    '''Adds an index on all columns in a table.
983
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
984
        index should be added on the first column.
985
        * If already_indexed, the pkey is assumed to have already been added
986
    '''
987
    cols = table_cols(db, table)
988
    if has_pkey:
989
        if has_pkey is not already_indexed: add_pkey(db, table)
990
        cols = cols[1:]
991
    for col in cols: add_index(db, col, table)
992

    
993
def add_row_num(db, table):
994
    '''Adds a row number column to a table. Its name is in row_num_col. It will
995
    be the primary key.'''
996
    table = sql_gen.as_Table(table).to_str(db)
997
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
998
        +' serial NOT NULL PRIMARY KEY', log_level=3)
999

    
1000
def drop_table(db, table):
1001
    table = sql_gen.as_Table(table)
1002
    return run_query(db, 'DROP TABLE IF EXISTS '+table.to_str(db)+' CASCADE')
1003

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

    
1030
def vacuum(db, table):
1031
    table = sql_gen.as_Table(table)
1032
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1033
        log_level=3))
1034

    
1035
def truncate(db, table, schema='public'):
1036
    table = sql_gen.as_Table(table, schema)
1037
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE')
1038

    
1039
def tables(db, schema_like='public', table_like='%', exact=False):
1040
    if exact: compare = '='
1041
    else: compare = 'LIKE'
1042
    
1043
    module = util.root_module(db.db)
1044
    if module == 'psycopg2':
1045
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1046
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1047
        return values(select(db, 'pg_tables', ['tablename'], conds,
1048
            order_by='tablename', log_level=4))
1049
    elif module == 'MySQLdb':
1050
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1051
            , cacheable=True, log_level=4))
1052
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1053

    
1054
def table_exists(db, table):
1055
    table = sql_gen.as_Table(table)
1056
    return list(tables(db, table.schema, table.name, exact=True)) != []
1057

    
1058
def errors_table(db, table, if_exists=True):
1059
    '''
1060
    @param if_exists If set, returns None if the errors table doesn't exist
1061
    @return None|sql_gen.Table
1062
    '''
1063
    table = sql_gen.as_Table(table)
1064
    if table.srcs != (): table = table.srcs[0]
1065
    
1066
    errors_table = sql_gen.suffixed_table(table, '.errors')
1067
    if if_exists and not table_exists(db, errors_table): return None
1068
    return errors_table
1069

    
1070
##### Database management
1071

    
1072
def empty_db(db, schema='public', **kw_args):
1073
    '''For kw_args, see tables()'''
1074
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1075

    
1076
##### Heuristic queries
1077

    
1078
def put(db, table, row, pkey_=None, row_ct_ref=None):
1079
    '''Recovers from errors.
1080
    Only works under PostgreSQL (uses INSERT RETURNING).
1081
    '''
1082
    row = sql_gen.ColDict(db, table, row)
1083
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
1084
    
1085
    try:
1086
        cur = insert(db, table, row, pkey_, recover=True)
1087
        if row_ct_ref != None and cur.rowcount >= 0:
1088
            row_ct_ref[0] += cur.rowcount
1089
        return value(cur)
1090
    except DuplicateKeyException, e:
1091
        row = sql_gen.ColDict(db, table,
1092
            util.dict_subset_right_join(row, e.cols))
1093
        return value(select(db, table, [pkey_], row, recover=True))
1094

    
1095
def get(db, table, row, pkey, row_ct_ref=None, create=False):
1096
    '''Recovers from errors'''
1097
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
1098
    except StopIteration:
1099
        if not create: raise
1100
        return put(db, table, row, pkey, row_ct_ref) # insert new row
1101

    
1102
def is_func_result(col):
1103
    return col.table.name.find('(') >= 0 and col.name == 'result'
1104

    
1105
def into_table_name(out_table, in_tables0, mapping, is_func):
1106
    def in_col_str(in_col):
1107
        in_col = sql_gen.remove_col_rename(in_col)
1108
        if isinstance(in_col, sql_gen.Col):
1109
            table = in_col.table
1110
            if table == in_tables0:
1111
                in_col = sql_gen.to_name_only_col(in_col)
1112
            elif is_func_result(in_col): in_col = table # omit col name
1113
        return str(in_col)
1114
    
1115
    str_ = str(out_table)
1116
    if is_func:
1117
        str_ += '('
1118
        
1119
        try: value_in_col = mapping['value']
1120
        except KeyError:
1121
            str_ += ', '.join((str(k)+'='+in_col_str(v)
1122
                for k, v in mapping.iteritems()))
1123
        else: str_ += in_col_str(value_in_col)
1124
        
1125
        str_ += ')'
1126
    else:
1127
        out_col = 'rank'
1128
        try: in_col = mapping[out_col]
1129
        except KeyError: str_ += '_pkeys'
1130
        else: # has a rank column, so hierarchical
1131
            str_ += '['+str(out_col)+'='+in_col_str(in_col)+']'
1132
    return str_
1133

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

    
1367
##### Data cleanup
1368

    
1369
def cleanup_table(db, table, cols):
1370
    table = sql_gen.as_Table(table)
1371
    cols = map(sql_gen.as_Col, cols)
1372
    
1373
    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '
1374
        +db.esc_value(r'\N')+')')
1375
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))
1376
        for v in cols]
1377
    
1378
    update(db, table, changes)
(24-24/36)