Project

General

Profile

1
# Database import/export
2

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

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

    
16
##### Exceptions
17

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

    
21
##### Data cleanup
22

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

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

    
52
##### Error tracking
53

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

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

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

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

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

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

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

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

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

    
222
##### Import
223

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

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

    
287
def put(db, table, row, pkey_=None, row_ct_ref=None):
288
    '''Recovers from errors.
289
    Only works under PostgreSQL (uses INSERT RETURNING).
290
    '''
291
    return put_table(db, table, [], row, row_ct_ref)
292

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

    
302
def is_func_result(col):
303
    return col.table.name.find('(') >= 0 and col.name == 'result'
304

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

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