Project

General

Profile

1 3077 aaronmk
# Database import/export
2
3 3534 aaronmk
import copy
4 3431 aaronmk
import operator
5 3714 aaronmk
import warnings
6 3431 aaronmk
7 3077 aaronmk
import exc
8
import dicts
9
import sql
10
import sql_gen
11
import strings
12
import util
13
14 3645 aaronmk
##### Exceptions
15
16
# Can't use built-in SyntaxError because it stringifies to only the first line
17
class SyntaxError(Exception): pass
18
19 3081 aaronmk
##### Data cleanup
20
21
def cleanup_table(db, table, cols):
22
    table = sql_gen.as_Table(table)
23
    cols = map(sql_gen.as_Col, cols)
24
25
    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '
26
        +db.esc_value(r'\N')+')')
27
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))
28
        for v in cols]
29
30
    sql.update(db, table, changes, in_place=True)
31
32 3078 aaronmk
##### Error tracking
33
34
def track_data_error(db, errors_table, cols, value, error_code, error):
35
    '''
36
    @param errors_table If None, does nothing.
37
    '''
38
    if errors_table == None or cols == (): return
39
40
    for col in cols:
41
        try:
42
            sql.insert(db, errors_table, dict(column=col.name, value=value,
43
                error_code=error_code, error=error), recover=True,
44
                cacheable=True, log_level=4)
45
        except sql.DuplicateKeyException: pass
46
47 3506 aaronmk
class ExcToErrorsTable(sql_gen.ExcToWarning):
48
    '''Handles an exception by saving it or converting it to a warning.'''
49 3511 aaronmk
    def __init__(self, return_, srcs, errors_table, value=None):
50 3506 aaronmk
        '''
51
        @param return_ See sql_gen.ExcToWarning
52
        @param srcs The column names for the errors table
53
        @param errors_table None|sql_gen.Table
54 3511 aaronmk
        @param value The value (or an expression for it) that caused the error
55 3506 aaronmk
        @pre The invalid value must be in a local variable "value" of type text.
56
        '''
57
        sql_gen.ExcToWarning.__init__(self, return_)
58
59 3511 aaronmk
        value = sql_gen.as_Code(value)
60
61 3506 aaronmk
        self.srcs = srcs
62
        self.errors_table = errors_table
63 3511 aaronmk
        self.value = value
64 3501 aaronmk
65 3506 aaronmk
    def to_str(self, db):
66
        if not self.srcs or self.errors_table == None:
67
            return sql_gen.ExcToWarning.to_str(self, db)
68
69 3459 aaronmk
        errors_table_cols = map(sql_gen.Col,
70
            ['column', 'value', 'error_code', 'error'])
71 3465 aaronmk
        col_names_query = sql.mk_select(db, sql_gen.NamedValues('c', None,
72 3506 aaronmk
            [[c.name] for c in self.srcs]), order_by=None)
73
        insert_query = sql.mk_insert_select(db, self.errors_table,
74
            errors_table_cols,
75 3465 aaronmk
            sql_gen.Values(errors_table_cols).to_str(db))+';\n'
76 3506 aaronmk
        return '''\
77 3459 aaronmk
-- Save error in errors table.
78
DECLARE
79
    error_code text := SQLSTATE;
80
    error text := SQLERRM;
81 3511 aaronmk
    value text := '''+self.value.to_str(db)+''';
82 3529 aaronmk
    "column" text;
83 3459 aaronmk
BEGIN
84
    -- Insert the value and error for *each* source column.
85 3529 aaronmk
'''+strings.indent(sql_gen.RowExcIgnore(None, col_names_query, insert_query,
86 3467 aaronmk
    row_var=errors_table_cols[0]).to_str(db))+'''
87 3459 aaronmk
END;
88 3501 aaronmk
89 3506 aaronmk
'''+self.return_.to_str(db)
90 3459 aaronmk
91 3507 aaronmk
def data_exception_handler(*args, **kw_args):
92 3506 aaronmk
    '''Handles a data_exception by saving it or converting it to a warning.
93
    For params, see ExcToErrorsTable().
94
    '''
95
    return sql_gen.data_exception_handler(ExcToErrorsTable(*args, **kw_args))
96
97 3078 aaronmk
def cast(db, type_, col, errors_table=None):
98
    '''Casts an (unrenamed) column or value.
99
    If errors_table set and col has srcs, saves errors in errors_table (using
100 3360 aaronmk
    col's srcs attr as source columns). Otherwise, converts errors to warnings.
101 3078 aaronmk
    @param col str|sql_gen.Col|sql_gen.Literal
102
    @param errors_table None|sql_gen.Table|str
103
    '''
104
    col = sql_gen.as_Col(col)
105
106 3112 aaronmk
    # Don't convert exceptions to warnings for user-supplied constants
107
    if isinstance(col, sql_gen.Literal): return sql_gen.Cast(type_, col)
108
109 3078 aaronmk
    assert not isinstance(col, sql_gen.NamedCol)
110
111 3460 aaronmk
    function_name = strings.first_word(type_)
112 3459 aaronmk
    srcs = col.srcs
113 3508 aaronmk
    save_errors = errors_table != None and srcs
114
    if save_errors: # function will be unique for the given srcs
115 3750 aaronmk
        function_name = strings.ustr(sql_gen.FunctionCall(function_name,
116 3508 aaronmk
            *map(sql_gen.to_name_only_col, srcs)))
117 3078 aaronmk
    function = db.TempFunction(function_name)
118
119 3464 aaronmk
    # Create function definition
120
    modifiers = 'STRICT'
121
    if not save_errors: modifiers = 'IMMUTABLE '+modifiers
122 3511 aaronmk
    value_param = sql_gen.FunctionParam('value', 'text')
123
    handler = data_exception_handler('RETURN NULL;\n', srcs, errors_table,
124
        value_param.name)
125 3464 aaronmk
    body = sql_gen.CustomCode(handler.to_str(db, '''\
126 3467 aaronmk
/* The explicit cast to the return type is needed to make the cast happen
127
inside the try block. (Implicit casts to the return type happen at the end
128
of the function, outside any block.) */
129
RETURN value::'''+type_+''';
130 3464 aaronmk
'''))
131
    body.lang='plpgsql'
132 3500 aaronmk
    sql.define_func(db, sql_gen.FunctionDef(function, type_, body,
133 3511 aaronmk
        [value_param], modifiers))
134 3464 aaronmk
135 3078 aaronmk
    return sql_gen.FunctionCall(function, col)
136
137 3538 aaronmk
def func_wrapper_exception_handler(db, return_, args, errors_table):
138 3524 aaronmk
    '''Handles a function call's data_exceptions.
139
    Supports PL/Python functions.
140
    @param return_ See data_exception_handler()
141
    @param args [arg...] Function call's args
142
    @param errors_table See data_exception_handler()
143
    '''
144
    args = filter(sql_gen.has_srcs, args)
145
146
    srcs = sql_gen.cross_join_srcs(args)
147 3538 aaronmk
    value = sql_gen.merge_not_null(db, ',', args)
148 3524 aaronmk
    return sql_gen.NestedExcHandler(
149
        data_exception_handler(return_, srcs, errors_table, value)
150
        , sql_gen.plpythonu_error_handler
151
        )
152
153 3078 aaronmk
def cast_temp_col(db, type_, col, errors_table=None):
154
    '''Like cast(), but creates a new column with the cast values if the input
155
    is a column.
156
    @return The new column or cast value
157
    '''
158
    def cast_(col): return cast(db, type_, col, errors_table)
159
160
    try: col = sql_gen.underlying_col(col)
161
    except sql_gen.NoUnderlyingTableException: return sql_gen.wrap(cast_, col)
162
163
    table = col.table
164 3173 aaronmk
    new_col = sql_gen.suffixed_col(col, '::'+strings.first_word(type_))
165 3078 aaronmk
    expr = cast_(col)
166
167
    # Add column
168
    new_typed_col = sql_gen.TypedCol(new_col.name, type_)
169 3750 aaronmk
    sql.add_col(db, table, new_typed_col, comment=strings.urepr(col)+'::'+type_)
170 3078 aaronmk
    new_col.name = new_typed_col.name # propagate any renaming
171
172 3110 aaronmk
    sql.update(db, table, [(new_col, expr)], in_place=True, recover=True)
173 3078 aaronmk
174
    return new_col
175
176
def errors_table(db, table, if_exists=True):
177
    '''
178
    @param if_exists If set, returns None if the errors table doesn't exist
179
    @return None|sql_gen.Table
180
    '''
181
    table = sql_gen.as_Table(table)
182
    if table.srcs != (): table = table.srcs[0]
183
184
    errors_table = sql_gen.suffixed_table(table, '.errors')
185
    if if_exists and not sql.table_exists(db, errors_table): return None
186
    return errors_table
187
188
##### Import
189
190 3077 aaronmk
def put(db, table, row, pkey_=None, row_ct_ref=None):
191
    '''Recovers from errors.
192
    Only works under PostgreSQL (uses INSERT RETURNING).
193
    '''
194 3633 aaronmk
    return put_table(db, table, [], row, row_ct_ref)
195 3077 aaronmk
196
def get(db, table, row, pkey, row_ct_ref=None, create=False):
197
    '''Recovers from errors'''
198
    try:
199
        return sql.value(sql.select(db, table, [pkey], row, limit=1,
200
            recover=True))
201
    except StopIteration:
202
        if not create: raise
203
        return put(db, table, row, pkey, row_ct_ref) # insert new row
204
205
def is_func_result(col):
206
    return col.table.name.find('(') >= 0 and col.name == 'result'
207
208
def into_table_name(out_table, in_tables0, mapping, is_func):
209
    def in_col_str(in_col):
210
        in_col = sql_gen.remove_col_rename(in_col)
211
        if isinstance(in_col, sql_gen.Col):
212
            table = in_col.table
213
            if table == in_tables0:
214
                in_col = sql_gen.to_name_only_col(in_col)
215
            elif is_func_result(in_col): in_col = table # omit col name
216 3750 aaronmk
        return strings.ustr(in_col)
217 3077 aaronmk
218
    str_ = str(out_table)
219
    if is_func:
220
        str_ += '('
221
222
        try: value_in_col = mapping['value']
223
        except KeyError:
224
            str_ += ', '.join((str(k)+'='+in_col_str(v)
225
                for k, v in mapping.iteritems()))
226
        else: str_ += in_col_str(value_in_col)
227
228
        str_ += ')'
229
    else:
230
        out_col = 'rank'
231
        try: in_col = mapping[out_col]
232
        except KeyError: str_ += '_pkeys'
233
        else: # has a rank column, so hierarchical
234
            str_ += '['+str(out_col)+'='+in_col_str(in_col)+']'
235
    return str_
236
237 3628 aaronmk
def put_table(db, out_table, in_tables, mapping, row_ct_ref=None, default=None,
238 3660 aaronmk
    col_defaults={}, on_error=exc.reraise):
239 3077 aaronmk
    '''Recovers from errors.
240
    Only works under PostgreSQL (uses INSERT RETURNING).
241
    IMPORTANT: Must be run at the *beginning* of a transaction.
242
    @param in_tables The main input table to select from, followed by a list of
243
        tables to join with it using the main input table's pkey
244
    @param mapping dict(out_table_col=in_table_col, ...)
245
        * out_table_col: str (*not* sql_gen.Col)
246
        * in_table_col: sql_gen.Col|literal-value
247
    @param default The *output* column to use as the pkey for missing rows.
248
        If this output column does not exist in the mapping, uses None.
249 3618 aaronmk
    @param col_defaults Default values for required columns.
250 3077 aaronmk
    @return sql_gen.Col Where the output pkeys are made available
251
    '''
252 3474 aaronmk
    import psycopg2.extensions
253
254 3077 aaronmk
    out_table = sql_gen.as_Table(out_table)
255
256
    def log_debug(msg): db.log_debug(msg, level=1.5)
257
    def col_ustr(str_):
258
        return strings.repr_no_u(sql_gen.remove_col_rename(str_))
259
260
    log_debug('********** New iteration **********')
261
    log_debug('Inserting these input columns into '+strings.as_tt(
262
        out_table.to_str(db))+':\n'+strings.as_table(mapping, ustr=col_ustr))
263
264
    is_function = sql.function_exists(db, out_table)
265
266
    if is_function: out_pkey = 'result'
267
    else: out_pkey = sql.pkey(db, out_table, recover=True)
268
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
269
270
    in_tables_ = in_tables[:] # don't modify input!
271 3432 aaronmk
    try: in_tables0 = in_tables_.pop(0) # first table is separate
272
    except IndexError: in_tables0 = None
273
    else:
274
        in_pkey = sql.pkey(db, in_tables0, recover=True)
275
        in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
276 3431 aaronmk
277
    # Determine if can use optimization for only literal values
278
    is_literals = not reduce(operator.or_, map(sql_gen.is_table_col,
279 3434 aaronmk
        mapping.values()), False)
280 3431 aaronmk
    is_literals_or_function = is_literals or is_function
281
282 3432 aaronmk
    if in_tables0 == None: errors_table_ = None
283
    else: errors_table_ = errors_table(db, in_tables0)
284 3431 aaronmk
285
    # Create input joins from list of input tables
286 3077 aaronmk
    input_joins = [in_tables0]+[sql_gen.Join(v,
287
        {in_pkey: sql_gen.join_same_not_null}) for v in in_tables_]
288
289 3433 aaronmk
    if mapping == {} and not is_function: # need >= one column for INSERT SELECT
290
        mapping = {out_pkey: None} # ColDict will replace with default value
291
292 3431 aaronmk
    if not is_literals:
293 3628 aaronmk
        into = sql_gen.as_Table(into_table_name(out_table, in_tables0, mapping,
294
            is_function))
295 3431 aaronmk
296
        # Set column sources
297
        in_cols = filter(sql_gen.is_table_col, mapping.values())
298
        for col in in_cols:
299
            if col.table == in_tables0: col.set_srcs(sql_gen.src_self)
300
301
        log_debug('Joining together input tables into temp table')
302
        # Place in new table so don't modify input and for speed
303
        in_table = sql_gen.Table('in')
304
        mapping = dicts.join(mapping, sql.flatten(db, in_table, input_joins,
305
            in_cols, preserve=[in_pkey_col]))
306
        input_joins = [in_table]
307
        db.log_debug('Temp table: '+strings.as_tt(in_table.to_str(db)), level=2)
308 3077 aaronmk
309 3692 aaronmk
    # Wrap mapping in a sql_gen.ColDict.
310
    # sql_gen.ColDict sanitizes both keys and values passed into it.
311
    # Do after applying dicts.join() because that returns a plain dict.
312 3077 aaronmk
    mapping = sql_gen.ColDict(db, out_table, mapping)
313
314
    # Resolve default value column
315
    if default != None:
316
        try: default = mapping[default]
317
        except KeyError:
318
            db.log_debug('Default value column '
319
                +strings.as_tt(strings.repr_no_u(default))
320
                +' does not exist in mapping, falling back to None', level=2.1)
321
            default = None
322
323 3287 aaronmk
    # Save default values for all rows since in_table may have rows deleted
324 3431 aaronmk
    if is_literals: pass
325
    elif is_function: full_in_table = in_table
326 3386 aaronmk
    else:
327
        full_in_table = sql_gen.suffixed_table(in_table, '_full')
328
        full_in_table_cols = [in_pkey_col]
329
        if default != None:
330
            full_in_table_cols.append(default)
331
            default = sql_gen.with_table(default, full_in_table)
332
        sql.run_query_into(db, sql.mk_select(db, in_table, full_in_table_cols,
333
            order_by=None), into=full_in_table, add_pkey_=True)
334 3287 aaronmk
335 3431 aaronmk
    if not is_literals:
336
        pkeys_names = [in_pkey, out_pkey]
337
        pkeys_cols = [in_pkey_col, out_pkey_col]
338 3077 aaronmk
339
    pkeys_table_exists_ref = [False]
340 3499 aaronmk
    def insert_into_pkeys(joins, cols=None, limit=None, **kw_args):
341 3477 aaronmk
        query = sql.mk_select(db, joins, cols, order_by=None, limit=limit)
342 3077 aaronmk
        if pkeys_table_exists_ref[0]:
343 3289 aaronmk
            sql.insert_select(db, into, pkeys_names, query, **kw_args)
344 3077 aaronmk
        else:
345 3304 aaronmk
            sql.run_query_into(db, query, into=into, add_pkey_=True, **kw_args)
346 3077 aaronmk
            pkeys_table_exists_ref[0] = True
347
348
    limit_ref = [None]
349 3418 aaronmk
    def mk_main_select(joins, cols):
350
        return sql.mk_select(db, joins, cols, limit=limit_ref[0], order_by=None)
351
352 3431 aaronmk
    if is_literals: insert_in_table = None
353
    else:
354
        insert_in_table = in_table
355
        insert_in_tables = [insert_in_table]
356 3352 aaronmk
    join_cols = sql_gen.ColDict(db, out_table)
357 3077 aaronmk
358
    exc_strs = set()
359
    def log_exc(e):
360
        e_str = exc.str_(e, first_line_only=True)
361
        log_debug('Caught exception: '+e_str)
362 3552 aaronmk
        if e_str in exc_strs: # avoid infinite loops
363
            log_debug('Exception already seen, handler broken')
364
            on_error(e)
365
            remove_all_rows()
366
        else: exc_strs.add(e_str)
367 3077 aaronmk
368
    def remove_all_rows():
369
        log_debug('Ignoring all rows')
370
        limit_ref[0] = 0 # just create an empty pkeys table
371
372 3352 aaronmk
    def ignore_cond(cond, e):
373 3704 aaronmk
        if is_literals: remove_all_rows()
374
        else:
375
            out_table_cols = sql_gen.ColDict(db, out_table)
376
            out_table_cols.update(util.dict_subset_right_join({},
377
                sql.table_cols(db, out_table)))
378
379
            in_cols = []
380
            cond = sql.map_expr(db, cond, mapping, in_cols)
381
            cond = sql.map_expr(db, cond, out_table_cols)
382
383
            track_data_error(db, errors_table_, sql_gen.cols_srcs(in_cols),
384
                None, e.cause.pgcode,
385
                strings.ensure_newl(e.cause.pgerror)+'condition: '+cond)
386
387
            not_cond = sql_gen.NotCond(sql_gen.CustomCode(cond))
388
            log_debug('Ignoring rows where '+strings.as_tt(not_cond.to_str(db)))
389
            sql.delete(db, insert_in_table, not_cond)
390 3352 aaronmk
391 3294 aaronmk
    not_null_cols = set()
392 3077 aaronmk
    def ignore(in_col, value, e):
393 3630 aaronmk
        if sql_gen.is_table_col(in_col):
394
            in_col = sql_gen.with_table(in_col, insert_in_table)
395
396
            track_data_error(db, errors_table_, in_col.srcs, value,
397
                e.cause.pgcode, e.cause.pgerror)
398
399
            sql.add_index(db, in_col, insert_in_table) # enable fast filtering
400
            if value != None and in_col not in not_null_cols:
401 3636 aaronmk
                log_debug('Replacing invalid value '+strings.as_tt(repr(value))
402 3637 aaronmk
                    +' with NULL in column '+strings.as_tt(in_col.to_str(db)))
403 3630 aaronmk
                sql.update(db, insert_in_table, [(in_col, None)],
404
                    sql_gen.ColValueCond(in_col, value))
405
            else:
406 3637 aaronmk
                log_debug('Ignoring rows with '+strings.as_tt(in_col.to_str(db))
407 3636 aaronmk
                    +' = '+strings.as_tt(repr(value)))
408 3630 aaronmk
                sql.delete(db, insert_in_table,
409
                    sql_gen.ColValueCond(in_col, value))
410
                if value == None: not_null_cols.add(in_col)
411 3293 aaronmk
        else:
412 3630 aaronmk
            assert isinstance(in_col, sql_gen.NamedCol)
413 3684 aaronmk
            in_value = sql_gen.remove_col_rename(in_col)
414
            assert sql_gen.is_literal(in_value)
415
            if value == in_value.value:
416
                if value != None:
417
                    log_debug('Replacing invalid literal '
418
                        +strings.as_tt(in_col.to_str(db))+' with NULL')
419
                    mapping[in_col.name] = None
420
                else:
421
                    remove_all_rows()
422
            # otherwise, all columns were being ignore()d because the specific
423
            # column couldn't be identified, and this was not the invalid column
424 3077 aaronmk
425 3431 aaronmk
    if not is_literals:
426
        def insert_pkeys_table(which):
427
            return sql_gen.Table(sql_gen.concat(in_table.name,
428
                '_insert_'+which+'_pkeys'))
429
        insert_out_pkeys = insert_pkeys_table('out')
430
        insert_in_pkeys = insert_pkeys_table('in')
431 3077 aaronmk
432 3918 aaronmk
    def mk_func_call():
433 3550 aaronmk
        args = dict(((k.name, v) for k, v in mapping.iteritems()))
434
        func_call = sql_gen.NamedCol(out_pkey,
435
            sql_gen.FunctionCall(out_table, **args))
436 3918 aaronmk
        return func_call, args
437 3550 aaronmk
438 3918 aaronmk
    if is_function and not is_literals:
439
        log_debug('Defining wrapper function')
440
441
        func_call, args = mk_func_call()
442
443
        # Create empty pkeys table so its row type can be used
444
        insert_into_pkeys(input_joins, [in_pkey_col, func_call], limit=0,
445
            recover=True)
446
        result_type = db.col_info(sql_gen.Col(out_pkey, into)).type
447
448
        ## Create error handling wrapper function
449
450
        wrapper = db.TempFunction(sql_gen.concat(into.name, '_wrap'))
451
452
        select_cols = [in_pkey_col]+args.values()
453
        row_var = copy.copy(sql_gen.row_var)
454
        row_var.set_srcs([in_table])
455
        in_pkey_var = sql_gen.Col(in_pkey, row_var)
456
457
        args = dict(((k, sql_gen.with_table(v, row_var))
458
            for k, v in args.iteritems()))
459
        func_call = sql_gen.FunctionCall(out_table, **args)
460
461
        def mk_return(result):
462
            return sql_gen.ReturnQuery(sql.mk_select(db,
463
                fields=[in_pkey_var, result], explain=False))
464
        exc_handler = func_wrapper_exception_handler(db,
465
            mk_return(sql_gen.Cast(result_type, None)), args.values(),
466
            errors_table_)
467
468
        sql.define_func(db, sql_gen.FunctionDef(wrapper, sql_gen.SetOf(into),
469
            sql_gen.RowExcIgnore(sql_gen.RowType(in_table),
470
                sql.mk_select(db, input_joins, order_by=None),
471
                mk_return(func_call), exc_handler=exc_handler)
472
            ))
473
        wrapper_table = sql_gen.FunctionCall(wrapper)
474 3550 aaronmk
475 3077 aaronmk
    # Do inserts and selects
476
    while True:
477 3473 aaronmk
        has_joins = join_cols != {}
478
479 3551 aaronmk
        # Handle unrecoverable errors in a special case
480
        if limit_ref[0] == 0:
481 3623 aaronmk
            if is_literals:
482 3686 aaronmk
                default = sql_gen.remove_col_rename(default)
483 3623 aaronmk
                log_debug('Returning default: '+strings.as_tt(repr(default)))
484
                return default
485 3551 aaronmk
            elif is_function: pass # empty pkeys table already created
486
            else:
487
                log_debug('Creating an empty output pkeys table')
488 3745 aaronmk
                has_joins = False # use the no-joins case
489 3551 aaronmk
                cur = sql.run_query_into(db, sql.mk_select(db, out_table,
490
                    [out_pkey], order_by=None, limit=0), into=insert_out_pkeys)
491
492 3077 aaronmk
            break # don't do main case
493
494
        # Prepare to insert new rows
495 3918 aaronmk
        if is_function:
496
            log_debug('Calling function on input rows')
497
            if is_literals: func_call, args = mk_func_call()
498 3077 aaronmk
        else:
499 3550 aaronmk
            log_debug('Trying to insert new rows')
500 3291 aaronmk
            insert_args = dict(recover=True, cacheable=False)
501
            if has_joins:
502
                insert_args.update(dict(ignore=True))
503
            else:
504 3431 aaronmk
                insert_args.update(dict(returning=out_pkey))
505
                if not is_literals:
506
                    insert_args.update(dict(into=insert_out_pkeys))
507 3291 aaronmk
            main_select = mk_main_select([insert_in_table], [sql_gen.with_table(
508
                c, insert_in_table) for c in mapping.values()])
509 3077 aaronmk
510 3292 aaronmk
        try:
511
            cur = None
512 3077 aaronmk
            if is_function:
513 3917 aaronmk
                if is_literals:
514
                    cur = sql.select(db, fields=[func_call], recover=True,
515
                        cacheable=True)
516 3499 aaronmk
                else: insert_into_pkeys(wrapper_table, recover=True)
517 3077 aaronmk
            else:
518 3292 aaronmk
                cur = sql.insert_select(db, out_table, mapping.keys(),
519 3077 aaronmk
                    main_select, **insert_args)
520
            break # insert successful
521
        except sql.MissingCastException, e:
522
            log_exc(e)
523
524
            out_col = e.col
525
            type_ = e.type
526
527
            log_debug('Casting '+strings.as_tt(out_col)+' input to '
528
                +strings.as_tt(type_))
529 3111 aaronmk
            in_col = mapping[out_col]
530
            while True:
531
                try:
532
                    mapping[out_col] = cast_temp_col(db, type_, in_col,
533
                        errors_table_)
534
                    break # cast successful
535
                except sql.InvalidValueException, e:
536
                    log_exc(e)
537
538 3294 aaronmk
                    ignore(in_col, e.value, e)
539 3077 aaronmk
        except sql.DuplicateKeyException, e:
540
            log_exc(e)
541
542 3274 aaronmk
            # Different rows violating different unique constraints not
543
            # supported
544
            assert not join_cols
545
546 3077 aaronmk
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
547
            log_debug('Ignoring existing rows, comparing on these columns:\n'
548
                +strings.as_inline_table(join_cols, ustr=col_ustr))
549 3102 aaronmk
550 3431 aaronmk
            if is_literals:
551
                return sql.value(sql.select(db, out_table, [out_pkey_col],
552
                    mapping, order_by=None))
553
554 3102 aaronmk
            # Uniquify input table to avoid internal duplicate keys
555
            insert_in_table = sql.distinct_table(db, insert_in_table,
556 3358 aaronmk
                join_cols.values())
557 3144 aaronmk
            insert_in_tables.append(insert_in_table)
558 3077 aaronmk
        except sql.NullValueException, e:
559
            log_exc(e)
560
561
            out_col, = e.cols
562
            try: in_col = mapping[out_col]
563 3618 aaronmk
            except KeyError, e:
564
                try: in_col = mapping[out_col] = col_defaults[out_col]
565
                except KeyError:
566
                    msg = 'Missing mapping for NOT NULL column '+out_col
567
                    log_debug(msg)
568 3713 aaronmk
                    if default == None: warnings.warn(UserWarning(msg))
569
                        # not an error because sometimes the mappings include
570
                        # extra tables which aren't used by the dataset
571 3618 aaronmk
                    remove_all_rows()
572 3294 aaronmk
            else: ignore(in_col, None, e)
573 3352 aaronmk
        except sql.CheckException, e:
574
            log_exc(e)
575
576
            ignore_cond(e.cond, e)
577 3413 aaronmk
        except sql.InvalidValueException, e:
578
            log_exc(e)
579
580
            for in_col in mapping.values(): ignore(in_col, e.value, e)
581 3474 aaronmk
        except psycopg2.extensions.TransactionRollbackError, e:
582
            log_exc(e)
583
            # retry
584 3077 aaronmk
        except sql.DatabaseErrors, e:
585
            log_exc(e)
586
587
            log_debug('No handler for exception')
588
            on_error(e)
589
            remove_all_rows()
590
        # after exception handled, rerun loop with additional constraints
591
592
    if cur != None and row_ct_ref != None and cur.rowcount >= 0:
593
        row_ct_ref[0] += cur.rowcount
594
595 3530 aaronmk
    if is_literals: return sql.value(cur)
596
597
    if is_function: pass # pkeys table already created
598 3077 aaronmk
    elif has_joins:
599
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
600
        log_debug('Getting output table pkeys of existing/inserted rows')
601 3306 aaronmk
        insert_into_pkeys(select_joins, pkeys_cols)
602 3077 aaronmk
    else:
603
        sql.add_row_num(db, insert_out_pkeys) # for joining with input pkeys
604
605
        log_debug('Getting input table pkeys of inserted rows')
606 3285 aaronmk
        # Note that mk_main_select() does not use ORDER BY. Instead, assume that
607
        # since the SELECT query is identical to the one used in INSERT SELECT,
608
        # its rows will be retrieved in the same order.
609 3077 aaronmk
        sql.run_query_into(db, mk_main_select(input_joins, [in_pkey]),
610
            into=insert_in_pkeys)
611
        sql.add_row_num(db, insert_in_pkeys) # for joining with output pkeys
612
613
        assert sql.table_row_count(db, insert_out_pkeys) == sql.table_row_count(
614
            db, insert_in_pkeys)
615
616
        log_debug('Combining output and input pkeys in inserted order')
617
        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,
618
            {sql.row_num_col: sql_gen.join_same_not_null})]
619
        insert_into_pkeys(pkey_joins, pkeys_names)
620
621
        sql.empty_temp(db, [insert_out_pkeys, insert_in_pkeys])
622
623 3531 aaronmk
    if limit_ref[0] == 0 or not is_function: # is_function doesn't leave holes
624 3187 aaronmk
        log_debug('Setting pkeys of missing rows to '
625
            +strings.as_tt(repr(default)))
626 3287 aaronmk
        missing_rows_joins = [full_in_table, sql_gen.Join(into,
627 3187 aaronmk
            {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
628
            # must use join_same_not_null or query will take forever
629
        insert_into_pkeys(missing_rows_joins,
630 3287 aaronmk
            [sql_gen.Col(in_pkey, full_in_table),
631
            sql_gen.NamedCol(out_pkey, default)])
632 3187 aaronmk
    # otherwise, there is already an entry for every row
633 3077 aaronmk
634 3530 aaronmk
    assert (sql.table_row_count(db, into)
635
        == sql.table_row_count(db, full_in_table))
636
637
    sql.empty_temp(db, insert_in_tables+[full_in_table])
638
639
    srcs = []
640 3619 aaronmk
    if is_function: srcs = sql_gen.cols_srcs(in_cols)
641 3530 aaronmk
    return sql_gen.Col(out_pkey, into, srcs)