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']
22

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

    
41
##### Error tracking
42

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

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

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

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

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

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

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

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

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

    
211
##### Import
212

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

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

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

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

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