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
    if not cols: return
29
    
30
    expr = 'trim(both from %s)'
31
    for null in null_strs: expr = 'nullif('+expr+', '+db.esc_value(null)+')'
32
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db))) for v in cols]
33
    
34
    while True:
35
        try:
36
            sql.update(db, table, changes, in_place=True, recover=True)
37
            break # successful
38
        except sql.NullValueException, e:
39
            db.log_debug('Caught exception: '+exc.str_(e))
40
            col, = e.cols
41
            sql.drop_not_null(db, col)
42

    
43
##### Error tracking
44

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

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

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

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

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

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

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

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

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

    
213
##### Import
214

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

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

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

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

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