Project

General

Profile

1
# Database access
2

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

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

    
18
##### Exceptions
19

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

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

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

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

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

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

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

    
63
class NameException(DbException): pass
64

    
65
class DuplicateKeyException(ConstraintException): pass
66

    
67
class NullValueException(ConstraintException): pass
68

    
69
class FunctionValueException(ExceptionWithNameValue): pass
70

    
71
class DuplicateTableException(ExceptionWithName): pass
72

    
73
class DuplicateFunctionException(ExceptionWithName): pass
74

    
75
class EmptyRowException(DbException): pass
76

    
77
##### Warnings
78

    
79
class DbWarning(UserWarning): pass
80

    
81
##### Result retrieval
82

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

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

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

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

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

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

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

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

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

    
108
##### Escaping
109

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

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

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

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

    
128
##### Database connections
129

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

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

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

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

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

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

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

    
400
connect = DbConn
401

    
402
##### Recoverable querying
403

    
404
def with_savepoint(db, func): return db.with_savepoint(func)
405

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

    
467
##### Basic queries
468

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

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

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

    
514
order_by_pkey = object() # tells mk_select() to order by the pkey
515

    
516
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
517

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    
850
##### Database structure queries
851

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

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

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

    
864
not_null_col = 'not_null_col'
865

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

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

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

    
929
row_num_col = '_row_num'
930

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

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

    
994
already_indexed = object() # tells add_indexes() the pkey has already been added
995

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

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

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

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

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

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

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

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

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

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

    
1099
##### Database management
1100

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

    
1105
##### Heuristic queries
1106

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

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

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

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

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

    
1402
##### Data cleanup
1403

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