Project

General

Profile

1
# Database access
2

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

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

    
18
##### Exceptions
19

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

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

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

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

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

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

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

    
63
class NameException(DbException): pass
64

    
65
class DuplicateKeyException(ConstraintException): pass
66

    
67
class NullValueException(ConstraintException): pass
68

    
69
class FunctionValueException(ExceptionWithNameValue): pass
70

    
71
class DuplicateTableException(ExceptionWithName): pass
72

    
73
class DuplicateFunctionException(ExceptionWithName): pass
74

    
75
class EmptyRowException(DbException): pass
76

    
77
##### Warnings
78

    
79
class DbWarning(UserWarning): pass
80

    
81
##### Result retrieval
82

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

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

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

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

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

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

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

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

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

    
108
##### Escaping
109

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

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

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

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

    
128
##### Database connections
129

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

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

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

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

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

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

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

    
377
connect = DbConn
378

    
379
##### Recoverable querying
380

    
381
def with_savepoint(db, func): return db.with_savepoint(func)
382

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

    
444
##### Basic queries
445

    
446
def next_version(name):
447
    version = 1 # first existing name was version 0
448
    match = re.match(r'^(.*)#(\d+)$', name)
449
    if match:
450
        name, version = match.groups()
451
        version = int(version)+1
452
    return sql_gen.add_suffix(name, '#'+str(version))
453

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

    
487
order_by_pkey = object() # tells mk_select() to order by the pkey
488

    
489
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
490

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

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

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

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

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

    
653
def insert(db, table, row, *args, **kw_args):
654
    '''For params, see insert_select()'''
655
    if lists.is_seq(row): cols = None
656
    else:
657
        cols = row.keys()
658
        row = row.values()
659
    row = list(row) # ensure that "== []" works
660
    
661
    if row == []: query = None
662
    else: query = sql_gen.Values(row).to_str(db)
663
    
664
    return insert_select(db, table, cols, query, *args, **kw_args)
665

    
666
def mk_update(db, table, changes=None, cond=None):
667
    '''
668
    @param changes [(col, new_value),...]
669
        * container can be any iterable type
670
        * col: sql_gen.Code|str (for col name)
671
        * new_value: sql_gen.Code|literal value
672
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
673
    @return str query
674
    '''
675
    query = 'UPDATE '+sql_gen.as_Table(table).to_str(db)+'\nSET\n'
676
    query += ',\n'.join((sql_gen.to_name_only_col(col, table).to_str(db)+' = '
677
        +sql_gen.as_Value(new_value).to_str(db) for col, new_value in changes))
678
    if cond != None: query += '\nWHERE\n'+cond.to_str(db)
679
    
680
    return query
681

    
682
def update(db, *args, **kw_args):
683
    '''For params, see mk_update() and run_query()'''
684
    recover = kw_args.pop('recover', None)
685
    
686
    return run_query(db, mk_update(db, *args, **kw_args), recover)
687

    
688
def last_insert_id(db):
689
    module = util.root_module(db.db)
690
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
691
    elif module == 'MySQLdb': return db.insert_id()
692
    else: return None
693

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

    
726
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
727
    '''For params, see mk_flatten_mapping()
728
    @return See return value of mk_flatten_mapping()
729
    '''
730
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
731
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
732
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
733
        into=into)
734
    return dict(items)
735

    
736
def mk_track_data_error(db, errors_table, cols, value, error_code, error):
737
    cols = map(sql_gen.to_name_only_col, cols)
738
    
739
    columns_cols = ['column']
740
    columns = sql_gen.NamedValues('columns', columns_cols,
741
        [[c.name] for c in cols])
742
    values_cols = ['value', 'error_code', 'error']
743
    values = sql_gen.NamedValues('values', values_cols,
744
        [value, error_code, error])
745
    
746
    select_cols = columns_cols+values_cols
747
    name_only_cols = map(sql_gen.to_name_only_col, select_cols)
748
    errors_table = sql_gen.NamedTable('errors', errors_table)
749
    joins = [columns, sql_gen.Join(values, type_='CROSS'),
750
        sql_gen.Join(errors_table, dict(zip(name_only_cols, select_cols)),
751
        sql_gen.filter_out)]
752
    
753
    return mk_insert_select(db, errors_table, name_only_cols,
754
        mk_select(db, joins, select_cols, order_by=None))
755

    
756
def track_data_error(db, errors_table, *args, **kw_args):
757
    '''
758
    @param errors_table If None, does nothing.
759
    '''
760
    if errors_table == None: return
761
    run_query(db, mk_track_data_error(db, errors_table, *args, **kw_args),
762
        cacheable=True, log_level=4)
763

    
764
def cast(db, type_, col, errors_table=None):
765
    '''Casts an (unrenamed) column or value.
766
    If errors_table set and col has srcs, saves errors in errors_table (using
767
    col's srcs attr as the source columns) and converts errors to warnings.
768
    @param col sql_gen.Col|sql_gen.Literal
769
    @param errors_table None|sql_gen.Table|str
770
    '''
771
    save_errors = (errors_table != None and isinstance(col, sql_gen.Col)
772
        and col.srcs != ())
773
    if not save_errors: # can't save errors
774
        return sql_gen.CustomCode(col.to_str(db)+'::'+type_) # just cast
775
    
776
    assert not isinstance(col, sql_gen.NamedCol)
777
    
778
    errors_table = sql_gen.as_Table(errors_table)
779
    srcs = map(sql_gen.to_name_only_col, col.srcs)
780
    function_name = str(sql_gen.FunctionCall(type_, *srcs))
781
    function = sql_gen.TempFunction(function_name, db.autocommit)
782
    
783
    while True:
784
        # Create function definition
785
        query = '''\
786
CREATE FUNCTION '''+function.to_str(db)+'''(value text)
787
RETURNS '''+type_+'''
788
LANGUAGE plpgsql
789
STRICT
790
AS $$
791
BEGIN
792
    /* The explicit cast to the return type is needed to make the cast happen
793
    inside the try block. (Implicit casts to the return type happen at the end
794
    of the function, outside any block.) */
795
    RETURN value::'''+type_+''';
796
EXCEPTION
797
    WHEN data_exception THEN
798
        -- Save error in errors table.
799
        -- Insert the value and error for *each* source column.
800
'''+mk_track_data_error(db, errors_table, srcs,
801
    *map(sql_gen.CustomCode, ['value', 'SQLSTATE', 'SQLERRM']))+''';
802
        
803
        RAISE WARNING '%', SQLERRM;
804
        RETURN NULL;
805
END;
806
$$;
807
'''
808
        
809
        # Create function
810
        try:
811
            run_query(db, query, recover=True, cacheable=True,
812
                log_ignore_excs=(DuplicateFunctionException,))
813
            break # successful
814
        except DuplicateFunctionException:
815
            function.name = next_version(function.name)
816
            # try again with next version of name
817
    
818
    return sql_gen.FunctionCall(function, col)
819

    
820
##### Database structure queries
821

    
822
def table_row_count(db, table, recover=None):
823
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
824
        order_by=None, start=0), recover=recover, log_level=3))
825

    
826
def table_cols(db, table, recover=None):
827
    return list(col_names(select(db, table, limit=0, order_by=None,
828
        recover=recover, log_level=4)))
829

    
830
def pkey(db, table, recover=None):
831
    '''Assumed to be first column in table'''
832
    return table_cols(db, table, recover)[0]
833

    
834
not_null_col = 'not_null_col'
835

    
836
def table_not_null_col(db, table, recover=None):
837
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
838
    if not_null_col in table_cols(db, table, recover): return not_null_col
839
    else: return pkey(db, table, recover)
840

    
841
def index_cols(db, table, index):
842
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
843
    automatically created. When you don't know whether something is a UNIQUE
844
    constraint or a UNIQUE index, use this function.'''
845
    module = util.root_module(db.db)
846
    if module == 'psycopg2':
847
        return list(values(run_query(db, '''\
848
SELECT attname
849
FROM
850
(
851
        SELECT attnum, attname
852
        FROM pg_index
853
        JOIN pg_class index ON index.oid = indexrelid
854
        JOIN pg_class table_ ON table_.oid = indrelid
855
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
856
        WHERE
857
            table_.relname = '''+db.esc_value(table)+'''
858
            AND index.relname = '''+db.esc_value(index)+'''
859
    UNION
860
        SELECT attnum, attname
861
        FROM
862
        (
863
            SELECT
864
                indrelid
865
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
866
                    AS indkey
867
            FROM pg_index
868
            JOIN pg_class index ON index.oid = indexrelid
869
            JOIN pg_class table_ ON table_.oid = indrelid
870
            WHERE
871
                table_.relname = '''+db.esc_value(table)+'''
872
                AND index.relname = '''+db.esc_value(index)+'''
873
        ) s
874
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
875
) s
876
ORDER BY attnum
877
'''
878
            , cacheable=True, log_level=4)))
879
    else: raise NotImplementedError("Can't list index columns for "+module+
880
        ' database')
881

    
882
def constraint_cols(db, table, constraint):
883
    module = util.root_module(db.db)
884
    if module == 'psycopg2':
885
        return list(values(run_query(db, '''\
886
SELECT attname
887
FROM pg_constraint
888
JOIN pg_class ON pg_class.oid = conrelid
889
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
890
WHERE
891
    relname = '''+db.esc_value(table)+'''
892
    AND conname = '''+db.esc_value(constraint)+'''
893
ORDER BY attnum
894
'''
895
            )))
896
    else: raise NotImplementedError("Can't list constraint columns for "+module+
897
        ' database')
898

    
899
row_num_col = '_row_num'
900

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

    
942
def add_pkey(db, table, cols=None, recover=None):
943
    '''Adds a primary key.
944
    @param cols [sql_gen.Col,...] The columns in the primary key.
945
        Defaults to the first column in the table.
946
    @pre The table must not already have a primary key.
947
    '''
948
    table = sql_gen.as_Table(table)
949
    if cols == None: cols = [pkey(db, table, recover)]
950
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
951
    
952
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
953
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
954
        log_ignore_excs=(DuplicateTableException,))
955

    
956
already_indexed = object() # tells add_indexes() the pkey has already been added
957

    
958
def add_indexes(db, table, has_pkey=True):
959
    '''Adds an index on all columns in a table.
960
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
961
        index should be added on the first column.
962
        * If already_indexed, the pkey is assumed to have already been added
963
    '''
964
    cols = table_cols(db, table)
965
    if has_pkey:
966
        if has_pkey is not already_indexed: add_pkey(db, table)
967
        cols = cols[1:]
968
    for col in cols: add_index(db, col, table)
969

    
970
def add_row_num(db, table):
971
    '''Adds a row number column to a table. Its name is in row_num_col. It will
972
    be the primary key.'''
973
    table = sql_gen.as_Table(table).to_str(db)
974
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
975
        +' serial NOT NULL PRIMARY KEY', log_level=3)
976

    
977
def create_table(db, table, cols, has_pkey=True, col_indexes=True):
978
    '''Creates a table.
979
    @param cols [sql_gen.TypedCol,...] The column names and types
980
    @param has_pkey If set, the first column becomes the primary key.
981
    @param col_indexes bool|[ref]
982
        * If True, indexes will be added on all non-pkey columns.
983
        * If a list reference, [0] will be set to a function to do this.
984
          This can be used to delay index creation until the table is populated.
985
    '''
986
    table = sql_gen.as_Table(table)
987
    
988
    if has_pkey:
989
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
990
        pkey.type += ' NOT NULL PRIMARY KEY'
991
    
992
    str_ = 'CREATE TABLE '+table.to_str(db)+' (\n'
993
    str_ += '\n, '.join(v.to_str(db) for v in cols)
994
    str_ += '\n);\n'
995
    run_query(db, str_, cacheable=True, log_level=2)
996
    
997
    # Add indexes
998
    if has_pkey: has_pkey = already_indexed
999
    def add_indexes_(): add_indexes(db, table, has_pkey)
1000
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1001
    elif col_indexes: add_indexes_() # add now
1002

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

    
1008
def truncate(db, table, schema='public'):
1009
    table = sql_gen.as_Table(table, schema)
1010
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE')
1011

    
1012
def tables(db, schema_like='public', table_like='%', exact=False):
1013
    if exact: compare = '='
1014
    else: compare = 'LIKE'
1015
    
1016
    module = util.root_module(db.db)
1017
    if module == 'psycopg2':
1018
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1019
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1020
        return values(select(db, 'pg_tables', ['tablename'], conds,
1021
            order_by='tablename', log_level=4))
1022
    elif module == 'MySQLdb':
1023
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1024
            , cacheable=True, log_level=4))
1025
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1026

    
1027
def table_exists(db, table):
1028
    table = sql_gen.as_Table(table)
1029
    return list(tables(db, table.schema, table.name, exact=True)) != []
1030

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

    
1043
##### Database management
1044

    
1045
def empty_db(db, schema='public', **kw_args):
1046
    '''For kw_args, see tables()'''
1047
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1048

    
1049
##### Heuristic queries
1050

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

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

    
1075
def is_func_result(col):
1076
    return col.table.name.find('(') >= 0 and col.name == 'result'
1077

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

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

    
1339
##### Data cleanup
1340

    
1341
def cleanup_table(db, table, cols):
1342
    table = sql_gen.as_Table(table)
1343
    cols = map(sql_gen.as_Col, cols)
1344
    
1345
    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '
1346
        +db.esc_value(r'\N')+')')
1347
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))
1348
        for v in cols]
1349
    
1350
    update(db, table, changes)
(24-24/36)