Project

General

Profile

1
# Database import/export
2

    
3
import copy
4
import operator
5
import warnings
6

    
7
import exc
8
import dicts
9
import sql
10
import sql_gen
11
import strings
12
import util
13

    
14
##### Exceptions
15

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

    
19
##### Data cleanup
20

    
21
null_strs = ['', '-', r'\N', 'NULL', 'UNKNOWN', 'nulo']
22

    
23
def cleanup_table(db, table):
24
    table = sql_gen.as_Table(table)
25
    cols = [sql_gen.as_Col(strings.ustr(c), table)
26
        for c in sql.table_cols(db, table)]
27
    cols = filter(lambda c: sql_gen.is_text_col(db, c), cols)
28
    
29
    expr = 'trim(both from %s)'
30
    for null in null_strs: expr = 'nullif('+expr+', '+db.esc_value(null)+')'
31
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db))) for v in cols]
32
    
33
    while True:
34
        try:
35
            sql.update(db, table, changes, in_place=True, recover=True)
36
            break # successful
37
        except sql.NullValueException, e:
38
            db.log_debug('Caught exception: '+exc.str_(e))
39
            col, = e.cols
40
            sql.drop_not_null(db, col)
41

    
42
##### Error tracking
43

    
44
def track_data_error(db, errors_table, cols, value, error_code, error):
45
    '''
46
    @param errors_table If None, does nothing.
47
    '''
48
    if errors_table == None or cols == (): return
49
    
50
    for col in cols:
51
        try:
52
            sql.insert(db, errors_table, dict(column=col.name, value=value,
53
                error_code=error_code, error=error), recover=True,
54
                cacheable=True, log_level=4)
55
        except sql.DuplicateKeyException: pass
56

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

    
99
'''+self.return_.to_str(db)
100

    
101
def data_exception_handler(*args, **kw_args):
102
    '''Handles a data_exception by saving it or converting it to a warning.
103
    For params, see ExcToErrorsTable().
104
    '''
105
    return sql_gen.data_exception_handler(ExcToErrorsTable(*args, **kw_args))
106

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

    
147
def func_wrapper_exception_handler(db, return_, args, errors_table):
148
    '''Handles a function call's data_exceptions.
149
    Supports PL/Python functions.
150
    @param return_ See data_exception_handler()
151
    @param args [arg...] Function call's args
152
    @param errors_table See data_exception_handler()
153
    '''
154
    args = filter(sql_gen.has_srcs, args)
155
    
156
    srcs = sql_gen.cross_join_srcs(args)
157
    value = sql_gen.merge_not_null(db, ',', args)
158
    return sql_gen.NestedExcHandler(
159
        data_exception_handler(return_, srcs, errors_table, value)
160
        , sql_gen.plpythonu_error_handler
161
        )
162

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

    
186
def errors_table(db, table, if_exists=True):
187
    '''
188
    @param if_exists If set, returns None if the errors table doesn't exist
189
    @return None|sql_gen.Table
190
    '''
191
    table = sql_gen.as_Table(table)
192
    if table.srcs != (): table = table.srcs[0]
193
    
194
    errors_table = sql_gen.suffixed_table(table, '.errors')
195
    if if_exists and not sql.table_exists(db, errors_table): return None
196
    return errors_table
197

    
198
def mk_errors_table(db, table):
199
    errors_table_ = errors_table(db, table, if_exists=False)
200
    if sql.table_exists(db, errors_table_, cacheable=False): return
201
    
202
    typed_cols = [
203
        sql_gen.TypedCol('column', 'text', nullable=False),
204
        sql_gen.TypedCol('value', 'text'),
205
        sql_gen.TypedCol('error_code', 'character varying(5)', nullable=False),
206
        sql_gen.TypedCol('error', 'text', nullable=False),
207
        ]
208
    sql.create_table(db, errors_table_, typed_cols, has_pkey=False)
209
    index_cols = ['column', 'value', 'error_code', 'error']
210
    sql.add_index(db, index_cols, errors_table_, unique=True)
211

    
212
##### Import
213

    
214
def put(db, table, row, pkey_=None, row_ct_ref=None):
215
    '''Recovers from errors.
216
    Only works under PostgreSQL (uses INSERT RETURNING).
217
    '''
218
    return put_table(db, table, [], row, row_ct_ref)
219

    
220
def get(db, table, row, pkey, row_ct_ref=None, create=False):
221
    '''Recovers from errors'''
222
    try:
223
        return sql.value(sql.select(db, table, [pkey], row, limit=1,
224
            recover=True))
225
    except StopIteration:
226
        if not create: raise
227
        return put(db, table, row, pkey, row_ct_ref) # insert new row
228

    
229
def is_func_result(col):
230
    return col.table.name.find('(') >= 0 and col.name == 'result'
231

    
232
def into_table_name(out_table, in_tables0, mapping, is_func):
233
    def in_col_str(in_col):
234
        in_col = sql_gen.remove_col_rename(in_col)
235
        if isinstance(in_col, sql_gen.Col):
236
            table = in_col.table
237
            if table == in_tables0:
238
                in_col = sql_gen.to_name_only_col(in_col)
239
            elif is_func_result(in_col): in_col = table # omit col name
240
        return strings.ustr(in_col)
241
    
242
    str_ = strings.ustr(out_table)
243
    if is_func:
244
        str_ += '('
245
        
246
        try: value_in_col = mapping['value']
247
        except KeyError:
248
            str_ += ', '.join((strings.ustr(k)+'='+in_col_str(v)
249
                for k, v in mapping.iteritems()))
250
        else: str_ += in_col_str(value_in_col)
251
        
252
        str_ += ')'
253
    else:
254
        out_col = 'rank'
255
        try: in_col = mapping[out_col]
256
        except KeyError: str_ += '_pkeys'
257
        else: # has a rank column, so hierarchical
258
            str_ += '['+strings.ustr(out_col)+'='+in_col_str(in_col)+']'
259
    return str_
260

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