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

    
1013
row_num_typed_col = sql_gen.TypedCol(row_num_col, 'serial', nullable=False,
1014
    constraints='PRIMARY KEY')
1015

    
1016
def add_row_num(db, table):
1017
    '''Adds a row number column to a table. Its name is in row_num_col. It will
1018
    be the primary key.'''
1019
    add_col(db, table, row_num_typed_col, log_level=3)
1020

    
1021
def cast_temp_col(db, type_, col, errors_table=None):
1022
    '''Like cast(), but creates a new column with the cast values if the input
1023
    is a column.
1024
    @return The new column or cast value
1025
    '''
1026
    def cast_(col): return cast(db, type_, col, errors_table)
1027
    
1028
    try: col = sql_gen.underlying_col(col)
1029
    except sql_gen.NoUnderlyingTableException: return sql_gen.wrap(cast_, col)
1030
    
1031
    table = col.table
1032
    new_col = sql_gen.Col(sql_gen.concat(col.name, '::'+type_), table, col.srcs)
1033
    expr = cast_(col)
1034
    add_col(db, table, sql_gen.TypedCol(new_col.name, type_))
1035
    update(db, table, [(new_col, expr)])
1036
    
1037
    return new_col
1038

    
1039
def drop_table(db, table):
1040
    table = sql_gen.as_Table(table)
1041
    return run_query(db, 'DROP TABLE IF EXISTS '+table.to_str(db)+' CASCADE')
1042

    
1043
def create_table(db, table, cols, has_pkey=True, col_indexes=True):
1044
    '''Creates a table.
1045
    @param cols [sql_gen.TypedCol,...] The column names and types
1046
    @param has_pkey If set, the first column becomes the primary key.
1047
    @param col_indexes bool|[ref]
1048
        * If True, indexes will be added on all non-pkey columns.
1049
        * If a list reference, [0] will be set to a function to do this.
1050
          This can be used to delay index creation until the table is populated.
1051
    '''
1052
    table = sql_gen.as_Table(table)
1053
    
1054
    if has_pkey:
1055
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1056
        pkey.constraints = 'PRIMARY KEY'
1057
    
1058
    str_ = 'CREATE TABLE '+table.to_str(db)+' (\n'
1059
    str_ += '\n, '.join(v.to_str(db) for v in cols)
1060
    str_ += '\n);\n'
1061
    run_query(db, str_, cacheable=True, log_level=2)
1062
    
1063
    # Add indexes
1064
    if has_pkey: has_pkey = already_indexed
1065
    def add_indexes_(): add_indexes(db, table, has_pkey)
1066
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1067
    elif col_indexes: add_indexes_() # add now
1068

    
1069
def vacuum(db, table):
1070
    table = sql_gen.as_Table(table)
1071
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1072
        log_level=3))
1073

    
1074
def truncate(db, table, schema='public'):
1075
    table = sql_gen.as_Table(table, schema)
1076
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE')
1077

    
1078
def tables(db, schema_like='public', table_like='%', exact=False):
1079
    if exact: compare = '='
1080
    else: compare = 'LIKE'
1081
    
1082
    module = util.root_module(db.db)
1083
    if module == 'psycopg2':
1084
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1085
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1086
        return values(select(db, 'pg_tables', ['tablename'], conds,
1087
            order_by='tablename', log_level=4))
1088
    elif module == 'MySQLdb':
1089
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1090
            , cacheable=True, log_level=4))
1091
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1092

    
1093
def table_exists(db, table):
1094
    table = sql_gen.as_Table(table)
1095
    return list(tables(db, table.schema, table.name, exact=True)) != []
1096

    
1097
def function_exists(db, function):
1098
    function = sql_gen.as_Function(function)
1099
    
1100
    info_table = sql_gen.Table('routines', 'information_schema')
1101
    conds = [('routine_name', function.name)]
1102
    schema = function.schema
1103
    if schema != None: conds.append(('routine_schema', schema))
1104
    # Exclude trigger functions, since they cannot be called directly
1105
    conds.append(('data_type', sql_gen.CompareCond('trigger', '!=')))
1106
    
1107
    return list(values(select(db, info_table, ['routine_name'], conds,
1108
        order_by='routine_schema', limit=1, log_level=4))) != []
1109
        # TODO: order_by search_path schema order
1110

    
1111
def errors_table(db, table, if_exists=True):
1112
    '''
1113
    @param if_exists If set, returns None if the errors table doesn't exist
1114
    @return None|sql_gen.Table
1115
    '''
1116
    table = sql_gen.as_Table(table)
1117
    if table.srcs != (): table = table.srcs[0]
1118
    
1119
    errors_table = sql_gen.suffixed_table(table, '.errors')
1120
    if if_exists and not table_exists(db, errors_table): return None
1121
    return errors_table
1122

    
1123
##### Database management
1124

    
1125
def empty_db(db, schema='public', **kw_args):
1126
    '''For kw_args, see tables()'''
1127
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1128

    
1129
##### Heuristic queries
1130

    
1131
def put(db, table, row, pkey_=None, row_ct_ref=None):
1132
    '''Recovers from errors.
1133
    Only works under PostgreSQL (uses INSERT RETURNING).
1134
    '''
1135
    row = sql_gen.ColDict(db, table, row)
1136
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
1137
    
1138
    try:
1139
        cur = insert(db, table, row, pkey_, recover=True)
1140
        if row_ct_ref != None and cur.rowcount >= 0:
1141
            row_ct_ref[0] += cur.rowcount
1142
        return value(cur)
1143
    except DuplicateKeyException, e:
1144
        row = sql_gen.ColDict(db, table,
1145
            util.dict_subset_right_join(row, e.cols))
1146
        return value(select(db, table, [pkey_], row, recover=True))
1147

    
1148
def get(db, table, row, pkey, row_ct_ref=None, create=False):
1149
    '''Recovers from errors'''
1150
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
1151
    except StopIteration:
1152
        if not create: raise
1153
        return put(db, table, row, pkey, row_ct_ref) # insert new row
1154

    
1155
def is_func_result(col):
1156
    return col.table.name.find('(') >= 0 and col.name == 'result'
1157

    
1158
def into_table_name(out_table, in_tables0, mapping, is_func):
1159
    def in_col_str(in_col):
1160
        in_col = sql_gen.remove_col_rename(in_col)
1161
        if isinstance(in_col, sql_gen.Col):
1162
            table = in_col.table
1163
            if table == in_tables0:
1164
                in_col = sql_gen.to_name_only_col(in_col)
1165
            elif is_func_result(in_col): in_col = table # omit col name
1166
        return str(in_col)
1167
    
1168
    str_ = str(out_table)
1169
    if is_func:
1170
        str_ += '('
1171
        
1172
        try: value_in_col = mapping['value']
1173
        except KeyError:
1174
            str_ += ', '.join((str(k)+'='+in_col_str(v)
1175
                for k, v in mapping.iteritems()))
1176
        else: str_ += in_col_str(value_in_col)
1177
        
1178
        str_ += ')'
1179
    else:
1180
        out_col = 'rank'
1181
        try: in_col = mapping[out_col]
1182
        except KeyError: str_ += '_pkeys'
1183
        else: # has a rank column, so hierarchical
1184
            str_ += '['+str(out_col)+'='+in_col_str(in_col)+']'
1185
    return str_
1186

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

    
1426
##### Data cleanup
1427

    
1428
def cleanup_table(db, table, cols):
1429
    table = sql_gen.as_Table(table)
1430
    cols = map(sql_gen.as_Col, cols)
1431
    
1432
    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '
1433
        +db.esc_value(r'\N')+')')
1434
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))
1435
        for v in cols]
1436
    
1437
    update(db, table, changes)
(24-24/36)