Project

General

Profile

1
# Database import/export
2

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

    
8
import csvs
9
import exc
10
import dicts
11
import sql
12
import sql_gen
13
import strings
14
import util
15

    
16
##### Exceptions
17

    
18
# Can't use built-in SyntaxError because it stringifies to only the first line
19
class SyntaxError(Exception): pass
20

    
21
##### Data cleanup
22

    
23
null_strs = ['', '-', r'\N', 'NULL', 'UNKNOWN', 'nulo']
24

    
25
def cleanup_table(db, table):
26
    table = sql_gen.as_Table(table)
27
    cols = filter(lambda c: sql_gen.is_text_col(db, c),
28
        sql.table_cols(db, table))
29
    try: cols.remove(sql.pkey_col(db, table))
30
    except ValueError: pass
31
    if not cols: return
32
    
33
    db.log_debug('Cleaning up table', level=1.5)
34
    
35
    expr = 'trim(both from %s)'
36
    for null in null_strs: expr = 'nullif('+expr+', '+db.esc_value(null)+')'
37
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db))) for v in cols]
38
    
39
    while True:
40
        try:
41
            sql.update(db, table, changes, in_place=True, recover=True)
42
            break # successful
43
        except sql.NullValueException, e:
44
            db.log_debug('Caught exception: '+exc.str_(e))
45
            col, = e.cols
46
            sql.drop_not_null(db, col)
47
    
48
    db.log_debug('Vacuuming and reanalyzing table', level=1.5)
49
    sql.vacuum(db, table)
50

    
51
##### Error tracking
52

    
53
def track_data_error(db, errors_table, cols, value, error_code, error):
54
    '''
55
    @param errors_table If None, does nothing.
56
    '''
57
    if errors_table == None or cols == (): return
58
    
59
    for col in cols:
60
        try:
61
            sql.insert(db, errors_table, dict(column=col.name, value=value,
62
                error_code=error_code, error=error), recover=True,
63
                cacheable=True, log_level=4)
64
        except sql.DuplicateKeyException: pass
65

    
66
class ExcToErrorsTable(sql_gen.ExcToWarning):
67
    '''Handles an exception by saving it or converting it to a warning.'''
68
    def __init__(self, return_, srcs, errors_table, value=None):
69
        '''
70
        @param return_ See sql_gen.ExcToWarning
71
        @param srcs The column names for the errors table
72
        @param errors_table None|sql_gen.Table
73
        @param value The value (or an expression for it) that caused the error
74
        @pre The invalid value must be in a local variable "value" of type text.
75
        '''
76
        sql_gen.ExcToWarning.__init__(self, return_)
77
        
78
        value = sql_gen.as_Code(value)
79
        
80
        self.srcs = srcs
81
        self.errors_table = errors_table
82
        self.value = value
83
    
84
    def to_str(self, db):
85
        if not self.srcs or self.errors_table == None:
86
            return sql_gen.ExcToWarning.to_str(self, db)
87
        
88
        errors_table_cols = map(sql_gen.Col,
89
            ['column', 'value', 'error_code', 'error'])
90
        col_names_query = sql.mk_select(db, sql_gen.NamedValues('c', None,
91
            [[c.name] for c in self.srcs]), order_by=None)
92
        insert_query = sql.mk_insert_select(db, self.errors_table,
93
            errors_table_cols,
94
            sql_gen.Values(errors_table_cols).to_str(db))+';\n'
95
        return '''\
96
-- Save error in errors table.
97
DECLARE
98
    error_code text := SQLSTATE;
99
    error text := SQLERRM;
100
    value text := '''+self.value.to_str(db)+''';
101
    "column" text;
102
BEGIN
103
    -- Insert the value and error for *each* source column.
104
'''+strings.indent(sql_gen.RowExcIgnore(None, col_names_query, insert_query,
105
    row_var=errors_table_cols[0]).to_str(db))+'''
106
END;
107

    
108
'''+self.return_.to_str(db)
109

    
110
def data_exception_handler(*args, **kw_args):
111
    '''Handles a data_exception by saving it or converting it to a warning.
112
    For params, see ExcToErrorsTable().
113
    '''
114
    return sql_gen.data_exception_handler(ExcToErrorsTable(*args, **kw_args))
115

    
116
def cast(db, type_, col, errors_table=None):
117
    '''Casts an (unrenamed) column or value.
118
    If errors_table set and col has srcs, saves errors in errors_table (using
119
    col's srcs attr as source columns). Otherwise, converts errors to warnings.
120
    @param col str|sql_gen.Col|sql_gen.Literal
121
    @param errors_table None|sql_gen.Table|str
122
    '''
123
    col = sql_gen.as_Col(col)
124
    
125
    # Don't convert exceptions to warnings for user-supplied constants
126
    if isinstance(col, sql_gen.Literal): return sql_gen.Cast(type_, col)
127
    
128
    assert not isinstance(col, sql_gen.NamedCol)
129
    
130
    function_name = strings.first_word(type_)
131
    srcs = col.srcs
132
    save_errors = errors_table != None and srcs
133
    if save_errors: # function will be unique for the given srcs
134
        function_name = strings.ustr(sql_gen.FunctionCall(function_name,
135
            *map(sql_gen.to_name_only_col, srcs)))
136
    function = db.TempFunction(function_name)
137
    
138
    # Create function definition
139
    modifiers = 'STRICT'
140
    if not save_errors: modifiers = 'IMMUTABLE '+modifiers
141
    value_param = sql_gen.FunctionParam('value', 'text')
142
    handler = data_exception_handler('RETURN NULL;\n', srcs, errors_table,
143
        value_param.name)
144
    body = sql_gen.CustomCode(handler.to_str(db, '''\
145
/* The explicit cast to the return type is needed to make the cast happen
146
inside the try block. (Implicit casts to the return type happen at the end
147
of the function, outside any block.) */
148
RETURN value::'''+type_+''';
149
'''))
150
    body.lang='plpgsql'
151
    sql.define_func(db, sql_gen.FunctionDef(function, type_, body,
152
        [value_param], modifiers))
153
    
154
    return sql_gen.FunctionCall(function, col)
155

    
156
def func_wrapper_exception_handler(db, return_, args, errors_table):
157
    '''Handles a function call's data_exceptions.
158
    Supports PL/Python functions.
159
    @param return_ See data_exception_handler()
160
    @param args [arg...] Function call's args
161
    @param errors_table See data_exception_handler()
162
    '''
163
    args = filter(sql_gen.has_srcs, args)
164
    
165
    srcs = sql_gen.cross_join_srcs(args)
166
    value = sql_gen.merge_not_null(db, ',', args)
167
    return sql_gen.NestedExcHandler(
168
        data_exception_handler(return_, srcs, errors_table, value)
169
        , sql_gen.plpythonu_error_handler
170
        )
171

    
172
def cast_temp_col(db, type_, col, errors_table=None):
173
    '''Like cast(), but creates a new column with the cast values if the input
174
    is a column.
175
    @return The new column or cast value
176
    '''
177
    def cast_(col): return cast(db, type_, col, errors_table)
178
    
179
    try: col = sql_gen.underlying_col(col)
180
    except sql_gen.NoUnderlyingTableException: return sql_gen.wrap(cast_, col)
181
    
182
    table = col.table
183
    new_col = sql_gen.suffixed_col(col, '::'+strings.first_word(type_))
184
    expr = cast_(col)
185
    
186
    # Add column
187
    new_typed_col = sql_gen.TypedCol(new_col.name, type_)
188
    sql.add_col(db, table, new_typed_col, comment=strings.urepr(col)+'::'+type_)
189
    new_col.name = new_typed_col.name # propagate any renaming
190
    
191
    sql.update(db, table, [(new_col, expr)], in_place=True, recover=True)
192
    
193
    return new_col
194

    
195
def errors_table(db, table, if_exists=True):
196
    '''
197
    @param if_exists If set, returns None if the errors table doesn't exist
198
    @return None|sql_gen.Table
199
    '''
200
    table = sql_gen.as_Table(table)
201
    if table.srcs != (): table = table.srcs[0]
202
    
203
    errors_table = sql_gen.suffixed_table(table, '.errors')
204
    if if_exists and not sql.table_exists(db, errors_table): return None
205
    return errors_table
206

    
207
def mk_errors_table(db, table):
208
    errors_table_ = errors_table(db, table, if_exists=False)
209
    if sql.table_exists(db, errors_table_, cacheable=False): return
210
    
211
    typed_cols = [
212
        sql_gen.TypedCol('column', 'text', nullable=False),
213
        sql_gen.TypedCol('value', 'text'),
214
        sql_gen.TypedCol('error_code', 'character varying(5)', nullable=False),
215
        sql_gen.TypedCol('error', 'text', nullable=False),
216
        ]
217
    sql.create_table(db, errors_table_, typed_cols, has_pkey=False)
218
    index_cols = ['column', 'value', 'error_code', 'error']
219
    sql.add_index(db, index_cols, errors_table_, unique=True)
220

    
221
##### Import
222

    
223
def append_csv(db, table, stream_info, stream, use_copy_from=True):
224
    assert sql.table_col_names(db, table) == stream_info.header
225
    
226
    def log(msg, level=1): db.log_debug(msg, level)
227
    
228
    dialect = stream_info.dialect
229
    if csvs.is_tsv(dialect): use_copy_from = False
230
    if use_copy_from:
231
        log('Using COPY FROM')
232
        
233
        # Create COPY FROM statement
234
        copy_from = ('COPY '+table.to_str(db)+' FROM STDIN DELIMITER '
235
            +db.esc_value(dialect.delimiter)+' NULL '+db.esc_value(''))
236
        assert not csvs.is_tsv(dialect)
237
        copy_from += ' CSV'
238
        if dialect.quoting != csv.QUOTE_NONE:
239
            quote_str = db.esc_value(dialect.quotechar)
240
            copy_from += ' QUOTE '+quote_str
241
            if dialect.doublequote: copy_from += ' ESCAPE '+quote_str
242
        copy_from += ';\n'
243
        
244
        log(copy_from, level=2)
245
        db.db.cursor().copy_expert(copy_from, stream)
246
    else:
247
        log('Using INSERT')
248
        cols_ct = len(stream_info.header)
249
        row_ct = 0
250
        inserted_row_ct = 0
251
        for row in csvs.make_reader(stream, dialect):
252
            row_ct += 1
253
            row = map(strings.to_unicode, row)
254
            util.list_set_length(row, cols_ct) # truncate extra cols
255
            if sql.insert(db, table, row, ignore=True, cacheable=False,
256
                log_level=5) != None:
257
                inserted_row_ct += 1
258
        log('Inserted '+str(inserted_row_ct)+' of '+str(row_ct)+' rows')
259

    
260
def import_csv(db, table, stream, use_copy_from=True):
261
    def log(msg, level=1): db.log_debug(msg, level)
262
    
263
    # Get format info
264
    info = csvs.stream_info(stream, parse_header=True)
265
    col_names = map(strings.to_unicode, info.header)
266
    for i, col in enumerate(col_names): # replace empty column names
267
        if col == '': col_names[i] = 'column_'+str(i)
268
    
269
    # Select schema and escape names
270
    def esc_name(name): return db.esc_name(name)
271
    
272
    typed_cols = [sql_gen.TypedCol(v, 'text') for v in col_names]
273
    
274
    log('Creating table')
275
    # Note that this is not rolled back if the import fails. Instead, it is
276
    # cached, and will not be re-run if the import is retried.
277
    sql.create_table(db, table, typed_cols, has_pkey=False, col_indexes=False)
278
    
279
    # Free memory used by deleted (rolled back) rows from any failed import.
280
    # This MUST be run so that the rows will be stored in inserted order, and
281
    # the row_num added after import will match up with the CSV's row order.
282
    sql.truncate(db, table)
283
    
284
    # Load the data
285
    def load(): append_csv(db, table, info, stream, use_copy_from)
286
    if use_copy_from: sql.with_savepoint(db, load)
287
    else: load()
288
    
289
    cleanup_table(db, table)
290

    
291
def put(db, table, row, pkey_=None, row_ct_ref=None):
292
    '''Recovers from errors.
293
    Only works under PostgreSQL (uses INSERT RETURNING).
294
    '''
295
    return put_table(db, table, [], row, row_ct_ref)
296

    
297
def get(db, table, row, pkey, row_ct_ref=None, create=False):
298
    '''Recovers from errors'''
299
    try:
300
        return sql.value(sql.select(db, table, [pkey], row, limit=1,
301
            recover=True))
302
    except StopIteration:
303
        if not create: raise
304
        return put(db, table, row, pkey, row_ct_ref) # insert new row
305

    
306
def is_func_result(col):
307
    return col.table.name.find('(') >= 0 and col.name == 'result'
308

    
309
def into_table_name(out_table, in_tables0, mapping, is_func):
310
    def in_col_str(in_col):
311
        in_col = sql_gen.remove_col_rename(in_col)
312
        if isinstance(in_col, sql_gen.Col):
313
            table = in_col.table
314
            if table == in_tables0:
315
                in_col = sql_gen.to_name_only_col(in_col)
316
            elif is_func_result(in_col): in_col = table # omit col name
317
        return strings.ustr(in_col)
318
    
319
    str_ = strings.ustr(out_table)
320
    if is_func:
321
        str_ += '('
322
        
323
        try: value_in_col = mapping['value']
324
        except KeyError:
325
            str_ += ', '.join((strings.ustr(k)+'='+in_col_str(v)
326
                for k, v in mapping.iteritems()))
327
        else: str_ += in_col_str(value_in_col)
328
        
329
        str_ += ')'
330
    else:
331
        out_col = 'rank'
332
        try: in_col = mapping[out_col]
333
        except KeyError: str_ += '_pkeys'
334
        else: # has a rank column, so hierarchical
335
            str_ += '['+strings.ustr(out_col)+'='+in_col_str(in_col)+']'
336
    return str_
337

    
338
def put_table(db, out_table, in_tables, mapping, row_ct_ref=None, default=None,
339
    col_defaults={}, on_error=exc.reraise):
340
    '''Recovers from errors.
341
    Only works under PostgreSQL (uses INSERT RETURNING).
342
    IMPORTANT: Must be run at the *beginning* of a transaction.
343
    @param in_tables The main input table to select from, followed by a list of
344
        tables to join with it using the main input table's pkey
345
    @param mapping dict(out_table_col=in_table_col, ...)
346
        * out_table_col: str (*not* sql_gen.Col)
347
        * in_table_col: sql_gen.Col|literal-value
348
    @param default The *output* column to use as the pkey for missing rows.
349
        If this output column does not exist in the mapping, uses None.
350
    @param col_defaults Default values for required columns.
351
    @return sql_gen.Col Where the output pkeys are made available
352
    '''
353
    import psycopg2.extensions
354
    
355
    out_table = sql_gen.as_Table(out_table)
356
    
357
    def log_debug(msg): db.log_debug(msg, level=1.5)
358
    def col_ustr(str_):
359
        return strings.repr_no_u(sql_gen.remove_col_rename(str_))
360
    
361
    log_debug('********** New iteration **********')
362
    log_debug('Inserting these input columns into '+strings.as_tt(
363
        out_table.to_str(db))+':\n'+strings.as_table(mapping, ustr=col_ustr))
364
    
365
    is_function = sql.function_exists(db, out_table)
366
    
367
    if is_function: row_ct_ref = None # only track inserted rows
368
    
369
    # Warn if inserting empty table rows
370
    if not mapping and not is_function: # functions with no args OK
371
        warnings.warn(UserWarning('Inserting empty table row(s)'))
372
    
373
    if is_function: out_pkey = 'result'
374
    else: out_pkey = sql.pkey_name(db, out_table, recover=True)
375
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
376
    
377
    in_tables_ = copy.copy(in_tables) # don't modify input!
378
    try: in_tables0 = in_tables_.pop(0) # first table is separate
379
    except IndexError: in_tables0 = None
380
    else:
381
        in_pkey = sql.pkey_name(db, in_tables0, recover=True)
382
        in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
383
    
384
    # Determine if can use optimization for only literal values
385
    is_literals = not reduce(operator.or_, map(sql_gen.is_table_col,
386
        mapping.values()), False)
387
    is_literals_or_function = is_literals or is_function
388
    
389
    if in_tables0 == None: errors_table_ = None
390
    else: errors_table_ = errors_table(db, in_tables0)
391
    
392
    # Create input joins from list of input tables
393
    input_joins = [in_tables0]+[sql_gen.Join(v,
394
        {in_pkey: sql_gen.join_same_not_null}) for v in in_tables_]
395
    
396
    orig_mapping = mapping.copy()
397
    if mapping == {} and not is_function: # need >= one column for INSERT SELECT
398
        mapping = {out_pkey: None} # ColDict will replace with default value
399
    
400
    if not is_literals:
401
        into = sql_gen.as_Table(into_table_name(out_table, in_tables0, mapping,
402
            is_function))
403
        # Ensure into's out_pkey is different from in_pkey by prepending table
404
        if is_function: into_out_pkey = out_pkey
405
        else: into_out_pkey = strings.ustr(out_pkey_col)
406
        
407
        # Set column sources
408
        in_cols = filter(sql_gen.is_table_col, mapping.values())
409
        for col in in_cols:
410
            if col.table == in_tables0: col.set_srcs(sql_gen.src_self)
411
        
412
        log_debug('Joining together input tables into temp table')
413
        # Place in new table so don't modify input and for speed
414
        in_table = sql_gen.Table('in')
415
        mapping = dicts.join(mapping, sql.flatten(db, in_table, input_joins,
416
            in_cols, preserve=[in_pkey_col]))
417
        input_joins = [in_table]
418
        db.log_debug('Temp table: '+strings.as_tt(in_table.to_str(db)), level=2)
419
    
420
    # Wrap mapping in a sql_gen.ColDict.
421
    # sql_gen.ColDict sanitizes both keys and values passed into it.
422
    # Do after applying dicts.join() because that returns a plain dict.
423
    mapping = sql_gen.ColDict(db, out_table, mapping)
424
    
425
    # Save all rows since in_table may have rows deleted
426
    if is_literals: pass
427
    elif is_function: full_in_table = in_table
428
    else:
429
        full_in_table = sql_gen.suffixed_table(in_table, '_full')
430
        sql.run_query_into(db, sql.mk_select(db, in_table, None, order_by=None),
431
            into=full_in_table, add_pkey_=True)
432
    
433
    pkeys_table_exists_ref = [False]
434
    def insert_into_pkeys(joins, cols=None, limit=None, **kw_args):
435
        query = sql.mk_select(db, joins, cols, order_by=None, limit=limit)
436
        if pkeys_table_exists_ref[0]:
437
            sql.insert_select(db, into, [in_pkey, into_out_pkey], query,
438
                **kw_args)
439
        else:
440
            sql.run_query_into(db, query, into=into, add_pkey_=True, **kw_args)
441
            pkeys_table_exists_ref[0] = True
442
    
443
    def mk_main_select(joins, cols):
444
        return sql.mk_select(db, joins, cols, order_by=None)
445
    
446
    if is_literals: insert_in_table = None
447
    else:
448
        insert_in_table = in_table
449
        insert_in_tables = [insert_in_table]
450
    join_cols = sql_gen.ColDict(db, out_table)
451
    
452
    exc_strs = set()
453
    def log_exc(e):
454
        e_str = exc.str_(e, first_line_only=True)
455
        log_debug('Caught exception: '+e_str)
456
        if e_str in exc_strs: # avoid infinite loops
457
            log_debug('Exception already seen, handler broken')
458
            on_error(e)
459
            remove_all_rows()
460
        else: exc_strs.add(e_str)
461
    
462
    ignore_all_ref = [False]
463
    def remove_all_rows():
464
        log_debug('Ignoring all rows')
465
        ignore_all_ref[0] = True # just return the default value column
466
    
467
    def ignore_cond(cond, e):
468
        if is_literals: remove_all_rows()
469
        else:
470
            out_table_cols = sql_gen.ColDict(db, out_table)
471
            out_table_cols.update(util.dict_subset_right_join({},
472
                sql.table_col_names(db, out_table)))
473
            
474
            in_cols = []
475
            cond = sql_gen.map_expr(db, cond, mapping, in_cols)
476
            cond = sql_gen.map_expr(db, cond, out_table_cols)
477
            
478
            track_data_error(db, errors_table_, sql_gen.cols_srcs(in_cols),
479
                None, e.cause.pgcode,
480
                strings.ensure_newl(e.cause.pgerror)+'condition: '+cond)
481
            
482
            not_cond = sql_gen.NotCond(sql_gen.CustomCode(cond))
483
            log_debug("Ignoring rows that don't satisfy "+strings.as_tt(cond))
484
            assert cond != sql_gen.true_expr # we know the constraint failed
485
            if cond == sql_gen.false_expr: remove_all_rows()
486
            else: sql.delete(db, insert_in_table, not_cond)
487
    
488
    not_null_cols = set()
489
    def ignore(in_col, value, e):
490
        if sql_gen.is_table_col(in_col):
491
            in_col = sql_gen.with_table(in_col, insert_in_table)
492
            
493
            track_data_error(db, errors_table_, in_col.srcs, value,
494
                e.cause.pgcode, e.cause.pgerror)
495
            
496
            sql.add_index(db, in_col, insert_in_table) # enable fast filtering
497
            if value != None and in_col not in not_null_cols:
498
                log_debug('Replacing invalid value '
499
                    +strings.as_tt(strings.urepr(value))+' with NULL in column '
500
                    +strings.as_tt(in_col.to_str(db)))
501
                sql.update(db, insert_in_table, [(in_col, None)],
502
                    sql_gen.ColValueCond(in_col, value))
503
            else:
504
                log_debug('Ignoring rows with '+strings.as_tt(in_col.to_str(db))
505
                    +' = '+strings.as_tt(strings.urepr(value)))
506
                sql.delete(db, insert_in_table,
507
                    sql_gen.ColValueCond(in_col, value))
508
                if value == None: not_null_cols.add(in_col)
509
        else:
510
            assert isinstance(in_col, sql_gen.NamedCol)
511
            in_value = sql_gen.remove_col_rename(in_col)
512
            assert sql_gen.is_literal(in_value)
513
            if value == in_value.value:
514
                if value != None:
515
                    log_debug('Replacing invalid literal '
516
                        +strings.as_tt(in_col.to_str(db))+' with NULL')
517
                    mapping[in_col.name] = None
518
                else:
519
                    remove_all_rows()
520
            # otherwise, all columns were being ignore()d because the specific
521
            # column couldn't be identified, and this was not the invalid column
522
    
523
    if not is_literals:
524
        def insert_pkeys_table(which):
525
            return sql_gen.Table(sql_gen.concat(in_table.name,
526
                '_insert_'+which+'_pkeys'))
527
        insert_out_pkeys = insert_pkeys_table('out')
528
        insert_in_pkeys = insert_pkeys_table('in')
529
    
530
    def mk_func_call():
531
        args = dict(((k.name, v) for k, v in mapping.iteritems()))
532
        return sql_gen.FunctionCall(out_table, **args), args
533
        
534
    if is_function and not is_literals:
535
        log_debug('Defining wrapper function')
536
        
537
        func_call, args = mk_func_call()
538
        func_call = sql_gen.NamedCol(into_out_pkey, func_call)
539
        
540
        # Create empty pkeys table so its row type can be used
541
        insert_into_pkeys(input_joins, [in_pkey_col, func_call], limit=0,
542
            recover=True)
543
        result_type = db.col_info(sql_gen.Col(into_out_pkey, into)).type
544
        
545
        ## Create error handling wrapper function
546
        
547
        wrapper = db.TempFunction(sql_gen.concat(into.name, '_wrap'))
548
        
549
        select_cols = [in_pkey_col]+args.values()
550
        row_var = copy.copy(sql_gen.row_var)
551
        row_var.set_srcs([in_table])
552
        in_pkey_var = sql_gen.Col(in_pkey, row_var)
553
        
554
        args = dict(((k, sql_gen.with_table(v, row_var))
555
            for k, v in args.iteritems()))
556
        func_call = sql_gen.FunctionCall(out_table, **args)
557
        
558
        def mk_return(result):
559
            return sql_gen.ReturnQuery(sql.mk_select(db,
560
                fields=[in_pkey_var, result], explain=False))
561
        exc_handler = func_wrapper_exception_handler(db,
562
            mk_return(sql_gen.Cast(result_type, None)), args.values(),
563
            errors_table_)
564
        
565
        sql.define_func(db, sql_gen.FunctionDef(wrapper, sql_gen.SetOf(into),
566
            sql_gen.RowExcIgnore(sql_gen.RowType(in_table),
567
                sql.mk_select(db, input_joins, order_by=None),
568
                mk_return(func_call), exc_handler=exc_handler)
569
            ))
570
        wrapper_table = sql_gen.FunctionCall(wrapper)
571
    
572
    missing_msg = None
573
    
574
    # Do inserts and selects
575
    while True:
576
        has_joins = join_cols != {}
577
        
578
        if ignore_all_ref[0]: break # unrecoverable error, so don't do main case
579
        
580
        # Prepare to insert new rows
581
        if is_function:
582
            log_debug('Calling function on input rows')
583
            if is_literals: func_call, args = mk_func_call()
584
        else:
585
            log_debug('Trying to insert new rows')
586
            insert_args = dict(recover=True, cacheable=False)
587
            if has_joins:
588
                insert_args.update(dict(ignore=True))
589
            else:
590
                insert_args.update(dict(returning=out_pkey))
591
                if not is_literals:
592
                    insert_args.update(dict(into=insert_out_pkeys))
593
            main_select = mk_main_select([insert_in_table], [sql_gen.with_table(
594
                c, insert_in_table) for c in mapping.values()])
595
        
596
        try:
597
            cur = None
598
            if is_function:
599
                if is_literals:
600
                    cur = sql.select(db, fields=[func_call], recover=True,
601
                        cacheable=True)
602
                else: insert_into_pkeys(wrapper_table, recover=True)
603
            else:
604
                cur = sql.insert_select(db, out_table, mapping.keys(),
605
                    main_select, **insert_args)
606
            break # insert successful
607
        except sql.MissingCastException, e:
608
            log_exc(e)
609
            
610
            type_ = e.type
611
            if e.col == None: out_cols = mapping.keys()
612
            else: out_cols = [e.col]
613
            
614
            for out_col in out_cols:
615
                log_debug('Casting '+strings.as_tt(strings.repr_no_u(out_col))
616
                    +' input to '+strings.as_tt(type_))
617
                in_col = mapping[out_col]
618
                while True:
619
                    try:
620
                        mapping[out_col] = cast_temp_col(db, type_, in_col,
621
                            errors_table_)
622
                        break # cast successful
623
                    except sql.InvalidValueException, e:
624
                        log_exc(e)
625
                        
626
                        ignore(in_col, e.value, e)
627
        except sql.DuplicateKeyException, e:
628
            log_exc(e)
629
            
630
            # Different rows violating different unique constraints not
631
            # supported
632
            assert not join_cols
633
            
634
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
635
            log_debug('Ignoring existing rows, comparing on these columns:\n'
636
                +strings.as_inline_table(join_cols, ustr=col_ustr))
637
            
638
            if is_literals:
639
                return sql.value(sql.select(db, out_table, [out_pkey_col],
640
                    join_cols, order_by=None))
641
            
642
            # Uniquify input table to avoid internal duplicate keys
643
            insert_in_table = sql.distinct_table(db, insert_in_table,
644
                join_cols.values())
645
            insert_in_tables.append(insert_in_table)
646
        except sql.NullValueException, e:
647
            log_exc(e)
648
            
649
            out_col, = e.cols
650
            try: in_col = mapping[out_col]
651
            except KeyError, e:
652
                try: in_col = mapping[out_col] = col_defaults[out_col]
653
                except KeyError:
654
                    missing_msg = 'Missing mapping for NOT NULL column '+out_col
655
                    log_debug(missing_msg)
656
                    remove_all_rows()
657
            else: ignore(in_col, None, e)
658
        except sql.CheckException, e:
659
            log_exc(e)
660
            
661
            ignore_cond(e.cond, e)
662
        except sql.InvalidValueException, e:
663
            log_exc(e)
664
            
665
            for in_col in mapping.values(): ignore(in_col, e.value, e)
666
        except psycopg2.extensions.TransactionRollbackError, e:
667
            log_exc(e)
668
            # retry
669
        except sql.DatabaseErrors, e:
670
            log_exc(e)
671
            
672
            log_debug('No handler for exception')
673
            on_error(e)
674
            remove_all_rows()
675
        # after exception handled, rerun loop with additional constraints
676
    
677
    # Resolve default value column
678
    if default != None:
679
        if ignore_all_ref[0]: mapping.update(orig_mapping) # use input cols
680
        try: default = mapping[default]
681
        except KeyError:
682
            db.log_debug('Default value column '
683
                +strings.as_tt(strings.repr_no_u(default))
684
                +' does not exist in mapping, falling back to None', level=2.1)
685
            default = None
686
        else: default = sql_gen.remove_col_rename(default)
687
    
688
    if missing_msg != None and default == None:
689
        warnings.warn(UserWarning(missing_msg))
690
        # not an error because sometimes the mappings include
691
        # extra tables which aren't used by the dataset
692
    
693
    # Handle unrecoverable errors
694
    if ignore_all_ref[0]:
695
        log_debug('Returning default: '+strings.as_tt(strings.urepr(default)))
696
        return default
697
    
698
    if cur != None and row_ct_ref != None and cur.rowcount >= 0:
699
        row_ct_ref[0] += cur.rowcount
700
    
701
    if is_literals: return sql.value(cur)
702
    
703
    if is_function: pass # pkeys table already created
704
    elif has_joins:
705
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
706
        log_debug('Getting output table pkeys of existing/inserted rows')
707
        insert_into_pkeys(select_joins, [in_pkey_col,
708
            sql_gen.NamedCol(into_out_pkey, out_pkey_col)])
709
    else:
710
        sql.add_row_num(db, insert_out_pkeys) # for joining with input pkeys
711
        
712
        log_debug('Getting input table pkeys of inserted rows')
713
        # Note that mk_main_select() does not use ORDER BY. Instead, assume that
714
        # since the SELECT query is identical to the one used in INSERT SELECT,
715
        # its rows will be retrieved in the same order.
716
        sql.run_query_into(db, mk_main_select(input_joins, [in_pkey]),
717
            into=insert_in_pkeys)
718
        sql.add_row_num(db, insert_in_pkeys) # for joining with output pkeys
719
        
720
        assert sql.table_row_count(db, insert_out_pkeys) == sql.table_row_count(
721
            db, insert_in_pkeys)
722
        
723
        log_debug('Combining output and input pkeys in inserted order')
724
        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,
725
            {sql.row_num_col: sql_gen.join_same_not_null})]
726
        in_col = sql_gen.Col(in_pkey, insert_in_pkeys)
727
        out_col = sql_gen.NamedCol(into_out_pkey,
728
            sql_gen.Col(out_pkey, insert_out_pkeys))
729
        insert_into_pkeys(pkey_joins, [in_col, out_col])
730
        
731
        sql.empty_temp(db, [insert_out_pkeys, insert_in_pkeys])
732
    
733
    if not is_function: # is_function doesn't leave holes
734
        log_debug('Setting pkeys of missing rows to '
735
            +strings.as_tt(strings.urepr(default)))
736
        if sql_gen.is_table_col(default):
737
            default = sql_gen.with_table(default, full_in_table)
738
        missing_rows_joins = [full_in_table, sql_gen.Join(into,
739
            {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
740
            # must use join_same_not_null or query will take forever
741
        insert_into_pkeys(missing_rows_joins,
742
            [sql_gen.Col(in_pkey, full_in_table),
743
            sql_gen.NamedCol(into_out_pkey, default)])
744
    # otherwise, there is already an entry for every row
745
    
746
    assert (sql.table_row_count(db, into)
747
        == sql.table_row_count(db, full_in_table))
748
    
749
    sql.empty_temp(db, insert_in_tables+[full_in_table])
750
    
751
    srcs = []
752
    if is_function: srcs = sql_gen.cols_srcs(in_cols)
753
    return sql_gen.Col(into_out_pkey, into, srcs)
(30-30/42)