Project

General

Profile

1 3077 aaronmk
# Database import/export
2
3
import exc
4
import dicts
5
import sql
6
import sql_gen
7
import strings
8
import util
9
10 3081 aaronmk
##### Data cleanup
11
12
def cleanup_table(db, table, cols):
13
    table = sql_gen.as_Table(table)
14
    cols = map(sql_gen.as_Col, cols)
15
16
    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '
17
        +db.esc_value(r'\N')+')')
18
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))
19
        for v in cols]
20
21
    sql.update(db, table, changes, in_place=True)
22
23 3078 aaronmk
##### Error tracking
24
25
def track_data_error(db, errors_table, cols, value, error_code, error):
26
    '''
27
    @param errors_table If None, does nothing.
28
    '''
29
    if errors_table == None or cols == (): return
30
31
    for col in cols:
32
        try:
33
            sql.insert(db, errors_table, dict(column=col.name, value=value,
34
                error_code=error_code, error=error), recover=True,
35
                cacheable=True, log_level=4)
36
        except sql.DuplicateKeyException: pass
37
38
def cast(db, type_, col, errors_table=None):
39
    '''Casts an (unrenamed) column or value.
40
    If errors_table set and col has srcs, saves errors in errors_table (using
41
    col's srcs attr as the source columns) and converts errors to warnings.
42
    @param col str|sql_gen.Col|sql_gen.Literal
43
    @param errors_table None|sql_gen.Table|str
44
    '''
45
    col = sql_gen.as_Col(col)
46
47 3112 aaronmk
    # Don't convert exceptions to warnings for user-supplied constants
48
    if isinstance(col, sql_gen.Literal): return sql_gen.Cast(type_, col)
49
50 3078 aaronmk
    assert not isinstance(col, sql_gen.NamedCol)
51
52 3112 aaronmk
    save_errors = (errors_table != None and isinstance(col, sql_gen.Col)
53
        and col.srcs != ())
54
    function_name = type_
55
    if save_errors:
56
        errors_table = sql_gen.as_Table(errors_table)
57
58
        srcs = map(sql_gen.to_name_only_col, col.srcs)
59
        function_name = str(sql_gen.FunctionCall(function_name, *srcs))
60 3078 aaronmk
    function = db.TempFunction(function_name)
61
62
    while True:
63
        # Create function definition
64
        query = '''\
65
CREATE FUNCTION '''+function.to_str(db)+'''(value text)
66
RETURNS '''+type_+'''
67
LANGUAGE plpgsql
68 3112 aaronmk
'''
69
        if not save_errors: query += 'IMMUTABLE '
70
        query += '''\
71 3078 aaronmk
STRICT
72
AS $$
73
BEGIN
74
    /* The explicit cast to the return type is needed to make the cast happen
75
    inside the try block. (Implicit casts to the return type happen at the end
76
    of the function, outside any block.) */
77
    RETURN value::'''+type_+''';
78
EXCEPTION
79
    WHEN data_exception THEN
80 3112 aaronmk
'''
81
        if save_errors:
82
            errors_table_cols = map(sql_gen.Col,
83
                ['column', 'value', 'error_code', 'error'])
84
            query += '''\
85 3078 aaronmk
        -- Save error in errors table.
86
        DECLARE
87
            error_code text := SQLSTATE;
88
            error text := SQLERRM;
89
            "column" text;
90
        BEGIN
91
            -- Insert the value and error for *each* source column.
92
            FOR "column" IN
93
'''+sql.mk_select(db, sql_gen.NamedValues('c', None, [[c.name] for c in srcs]),
94
    order_by=None, start=0)+'''
95
            LOOP
96
                BEGIN
97
'''+sql.mk_insert_select(db, errors_table, errors_table_cols,
98
    sql_gen.Values(errors_table_cols).to_str(db))+''';
99
                EXCEPTION
100
                    WHEN unique_violation THEN NULL; -- continue to next row
101
                END;
102
            END LOOP;
103
        END;
104
105 3112 aaronmk
'''
106
        query += '''\
107 3078 aaronmk
        RAISE WARNING '%', SQLERRM;
108
        RETURN NULL;
109
END;
110
$$;
111
'''
112
113
        # Create function
114
        try:
115
            sql.run_query(db, query, recover=True, cacheable=True,
116
                log_ignore_excs=(sql.DuplicateException,))
117
            break # successful
118
        except sql.DuplicateException:
119
            function.name = sql.next_version(function.name)
120
            # try again with next version of name
121
122
    return sql_gen.FunctionCall(function, col)
123
124
def cast_temp_col(db, type_, col, errors_table=None):
125
    '''Like cast(), but creates a new column with the cast values if the input
126
    is a column.
127
    @return The new column or cast value
128
    '''
129
    def cast_(col): return cast(db, type_, col, errors_table)
130
131
    try: col = sql_gen.underlying_col(col)
132
    except sql_gen.NoUnderlyingTableException: return sql_gen.wrap(cast_, col)
133
134
    table = col.table
135
    new_col = sql_gen.Col(sql_gen.concat(col.name, '::'+type_), table, col.srcs)
136
    expr = cast_(col)
137
138
    # Add column
139
    new_typed_col = sql_gen.TypedCol(new_col.name, type_)
140
    sql.add_col(db, table, new_typed_col, comment='src: '+repr(col))
141
    new_col.name = new_typed_col.name # propagate any renaming
142
143 3110 aaronmk
    sql.update(db, table, [(new_col, expr)], in_place=True, recover=True)
144 3078 aaronmk
    sql.add_index(db, new_col)
145
146
    return new_col
147
148
def errors_table(db, table, if_exists=True):
149
    '''
150
    @param if_exists If set, returns None if the errors table doesn't exist
151
    @return None|sql_gen.Table
152
    '''
153
    table = sql_gen.as_Table(table)
154
    if table.srcs != (): table = table.srcs[0]
155
156
    errors_table = sql_gen.suffixed_table(table, '.errors')
157
    if if_exists and not sql.table_exists(db, errors_table): return None
158
    return errors_table
159
160
##### Import
161
162 3077 aaronmk
def put(db, table, row, pkey_=None, row_ct_ref=None):
163
    '''Recovers from errors.
164
    Only works under PostgreSQL (uses INSERT RETURNING).
165
    '''
166
    row = sql_gen.ColDict(db, table, row)
167
    if pkey_ == None: pkey_ = sql.pkey(db, table, recover=True)
168
169
    try:
170 3131 aaronmk
        cur = sql.insert(db, table, row, pkey_, recover=True, log_level=3.5)
171 3077 aaronmk
        if row_ct_ref != None and cur.rowcount >= 0:
172
            row_ct_ref[0] += cur.rowcount
173
        return sql.value(cur)
174
    except sql.DuplicateKeyException, e:
175
        row = sql_gen.ColDict(db, table,
176
            util.dict_subset_right_join(row, e.cols))
177 3131 aaronmk
        return sql.value(sql.select(db, table, [pkey_], row, recover=True,
178
            log_level=3.5))
179 3077 aaronmk
180
def get(db, table, row, pkey, row_ct_ref=None, create=False):
181
    '''Recovers from errors'''
182
    try:
183
        return sql.value(sql.select(db, table, [pkey], row, limit=1,
184
            recover=True))
185
    except StopIteration:
186
        if not create: raise
187
        return put(db, table, row, pkey, row_ct_ref) # insert new row
188
189
def is_func_result(col):
190
    return col.table.name.find('(') >= 0 and col.name == 'result'
191
192
def into_table_name(out_table, in_tables0, mapping, is_func):
193
    def in_col_str(in_col):
194
        in_col = sql_gen.remove_col_rename(in_col)
195
        if isinstance(in_col, sql_gen.Col):
196
            table = in_col.table
197
            if table == in_tables0:
198
                in_col = sql_gen.to_name_only_col(in_col)
199
            elif is_func_result(in_col): in_col = table # omit col name
200
        return str(in_col)
201
202
    str_ = str(out_table)
203
    if is_func:
204
        str_ += '('
205
206
        try: value_in_col = mapping['value']
207
        except KeyError:
208
            str_ += ', '.join((str(k)+'='+in_col_str(v)
209
                for k, v in mapping.iteritems()))
210
        else: str_ += in_col_str(value_in_col)
211
212
        str_ += ')'
213
    else:
214
        out_col = 'rank'
215
        try: in_col = mapping[out_col]
216
        except KeyError: str_ += '_pkeys'
217
        else: # has a rank column, so hierarchical
218
            str_ += '['+str(out_col)+'='+in_col_str(in_col)+']'
219
    return str_
220
221
def put_table(db, out_table, in_tables, mapping, row_ct_ref=None, into=None,
222
    default=None, is_func=False, on_error=exc.raise_):
223
    '''Recovers from errors.
224
    Only works under PostgreSQL (uses INSERT RETURNING).
225
    IMPORTANT: Must be run at the *beginning* of a transaction.
226
    @param in_tables The main input table to select from, followed by a list of
227
        tables to join with it using the main input table's pkey
228
    @param mapping dict(out_table_col=in_table_col, ...)
229
        * out_table_col: str (*not* sql_gen.Col)
230
        * in_table_col: sql_gen.Col|literal-value
231
    @param into The table to contain the output and input pkeys.
232
        Defaults to `out_table.name+'_pkeys'`.
233
    @param default The *output* column to use as the pkey for missing rows.
234
        If this output column does not exist in the mapping, uses None.
235
    @param is_func Whether out_table is the name of a SQL function, not a table
236
    @return sql_gen.Col Where the output pkeys are made available
237
    '''
238
    out_table = sql_gen.as_Table(out_table)
239
240
    def log_debug(msg): db.log_debug(msg, level=1.5)
241
    def col_ustr(str_):
242
        return strings.repr_no_u(sql_gen.remove_col_rename(str_))
243
244
    log_debug('********** New iteration **********')
245
    log_debug('Inserting these input columns into '+strings.as_tt(
246
        out_table.to_str(db))+':\n'+strings.as_table(mapping, ustr=col_ustr))
247
248
    is_function = sql.function_exists(db, out_table)
249
250
    if is_function: out_pkey = 'result'
251
    else: out_pkey = sql.pkey(db, out_table, recover=True)
252
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
253
254
    if mapping == {}: # need at least one column for INSERT SELECT
255
        mapping = {out_pkey: None} # ColDict will replace with default value
256
257
    # Create input joins from list of input tables
258
    in_tables_ = in_tables[:] # don't modify input!
259
    in_tables0 = in_tables_.pop(0) # first table is separate
260 3078 aaronmk
    errors_table_ = errors_table(db, in_tables0)
261 3077 aaronmk
    in_pkey = sql.pkey(db, in_tables0, recover=True)
262
    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
263
    input_joins = [in_tables0]+[sql_gen.Join(v,
264
        {in_pkey: sql_gen.join_same_not_null}) for v in in_tables_]
265
266
    if into == None:
267
        into = into_table_name(out_table, in_tables0, mapping, is_func)
268
    into = sql_gen.as_Table(into)
269
270
    # Set column sources
271
    in_cols = filter(sql_gen.is_table_col, mapping.values())
272
    for col in in_cols:
273
        if col.table == in_tables0: col.set_srcs(sql_gen.src_self)
274
275
    log_debug('Joining together input tables into temp table')
276
    # Place in new table for speed and so don't modify input if values edited
277
    in_table = sql_gen.Table('in')
278
    mapping = dicts.join(mapping, sql.flatten(db, in_table, input_joins,
279
        in_cols, preserve=[in_pkey_col], start=0))
280
    input_joins = [in_table]
281
    db.log_debug('Temp table: '+strings.as_tt(in_table.to_str(db)), level=2)
282
283
    mapping = sql_gen.ColDict(db, out_table, mapping)
284
        # after applying dicts.join() because that returns a plain dict
285
286
    # Resolve default value column
287
    if default != None:
288
        try: default = mapping[default]
289
        except KeyError:
290
            db.log_debug('Default value column '
291
                +strings.as_tt(strings.repr_no_u(default))
292
                +' does not exist in mapping, falling back to None', level=2.1)
293
            default = None
294
295
    pkeys_names = [in_pkey, out_pkey]
296
    pkeys_cols = [in_pkey_col, out_pkey_col]
297
298
    pkeys_table_exists_ref = [False]
299
    def insert_into_pkeys(joins, cols, distinct=False):
300
        kw_args = {}
301
        if distinct: kw_args.update(dict(distinct_on=[in_pkey_col]))
302
        query = sql.mk_select(db, joins, cols, order_by=None, start=0,
303
            **kw_args)
304
305
        if pkeys_table_exists_ref[0]:
306
            sql.insert_select(db, into, pkeys_names, query)
307
        else:
308
            sql.run_query_into(db, query, into=into)
309
            pkeys_table_exists_ref[0] = True
310
311
    limit_ref = [None]
312 3106 aaronmk
    insert_in_table = in_table
313 3144 aaronmk
    insert_in_tables = [insert_in_table]
314 3077 aaronmk
    conds = set()
315
    distinct_on = sql_gen.ColDict(db, out_table)
316
    def mk_main_select(joins, cols):
317 3106 aaronmk
        conds_ = [(sql_gen.with_table(k, insert_in_table), v) for k, v in conds]
318
        return sql.mk_select(db, joins, cols, conds_, limit=limit_ref[0],
319 3102 aaronmk
            start=0)
320 3077 aaronmk
321
    exc_strs = set()
322
    def log_exc(e):
323
        e_str = exc.str_(e, first_line_only=True)
324
        log_debug('Caught exception: '+e_str)
325
        assert e_str not in exc_strs # avoid infinite loops
326
        exc_strs.add(e_str)
327
328
    def remove_all_rows():
329
        log_debug('Ignoring all rows')
330
        limit_ref[0] = 0 # just create an empty pkeys table
331
332
    def ignore(in_col, value, e):
333 3078 aaronmk
        track_data_error(db, errors_table_, in_col.srcs, value,
334 3077 aaronmk
            e.cause.pgcode, e.cause.pgerror)
335
        log_debug('Ignoring rows with '+strings.as_tt(repr(in_col))+' = '
336
            +strings.as_tt(repr(value)))
337
338
    def remove_rows(in_col, value, e):
339
        ignore(in_col, value, e)
340
        cond = (in_col, sql_gen.CompareCond(value, '!='))
341
        assert cond not in conds # avoid infinite loops
342
        conds.add(cond)
343
344
    def invalid2null(in_col, value, e):
345
        ignore(in_col, value, e)
346
        sql.update(db, in_table, [(in_col, None)],
347
            sql_gen.ColValueCond(in_col, value))
348
349
    def insert_pkeys_table(which):
350
        return sql_gen.Table(sql_gen.concat(in_table.name,
351
            '_insert_'+which+'_pkeys'))
352
    insert_out_pkeys = insert_pkeys_table('out')
353
    insert_in_pkeys = insert_pkeys_table('in')
354
355
    # Do inserts and selects
356
    join_cols = sql_gen.ColDict(db, out_table)
357
    while True:
358
        if limit_ref[0] == 0: # special case
359
            log_debug('Creating an empty pkeys table')
360
            cur = sql.run_query_into(db, sql.mk_select(db, out_table,
361
                [out_pkey], limit=limit_ref[0]), into=insert_out_pkeys)
362
            break # don't do main case
363
364
        has_joins = join_cols != {}
365
366
        log_debug('Trying to insert new rows')
367
368
        # Prepare to insert new rows
369
        insert_args = dict(recover=True, cacheable=False)
370
        if has_joins:
371
            insert_args.update(dict(ignore=True))
372
        else:
373
            insert_args.update(dict(returning=out_pkey, into=insert_out_pkeys))
374 3102 aaronmk
        main_select = mk_main_select([insert_in_table],
375
            [sql_gen.with_table(c, insert_in_table) for c in mapping.values()])
376 3077 aaronmk
377
        def main_insert():
378
            if is_function:
379
                log_debug('Calling function on input rows')
380
                args = dict(((k.name, v) for k, v in mapping.iteritems()))
381
                func_call = sql_gen.NamedCol(out_pkey,
382
                    sql_gen.FunctionCall(out_table, **args))
383
                insert_into_pkeys(input_joins, [in_pkey_col, func_call])
384
                return None
385
            else:
386
                return sql.insert_select(db, out_table, mapping.keys(),
387
                    main_select, **insert_args)
388
389
        try:
390
            cur = sql.with_savepoint(db, main_insert)
391
            break # insert successful
392
        except sql.MissingCastException, e:
393
            log_exc(e)
394
395
            out_col = e.col
396
            type_ = e.type
397
398
            log_debug('Casting '+strings.as_tt(out_col)+' input to '
399
                +strings.as_tt(type_))
400 3111 aaronmk
            in_col = mapping[out_col]
401
            while True:
402
                try:
403
                    mapping[out_col] = cast_temp_col(db, type_, in_col,
404
                        errors_table_)
405
                    break # cast successful
406
                except sql.InvalidValueException, e:
407
                    log_exc(e)
408
409
                    invalid2null(in_col, e.value, e)
410 3077 aaronmk
        except sql.DuplicateKeyException, e:
411
            log_exc(e)
412
413
            old_join_cols = join_cols.copy()
414
            distinct_on.update(util.dict_subset(mapping, e.cols))
415
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
416
            log_debug('Ignoring existing rows, comparing on these columns:\n'
417
                +strings.as_inline_table(join_cols, ustr=col_ustr))
418
            assert join_cols != old_join_cols # avoid infinite loops
419 3102 aaronmk
420
            # Uniquify input table to avoid internal duplicate keys
421
            insert_in_table = sql.distinct_table(db, insert_in_table,
422
                filter(sql_gen.is_table_col, distinct_on.values()))
423 3144 aaronmk
            insert_in_tables.append(insert_in_table)
424 3077 aaronmk
        except sql.NullValueException, e:
425
            log_exc(e)
426
427
            out_col, = e.cols
428
            try: in_col = mapping[out_col]
429
            except KeyError:
430
                log_debug('Missing mapping for NOT NULL column '+out_col)
431
                remove_all_rows()
432
            else: remove_rows(in_col, None, e)
433
        except sql.DatabaseErrors, e:
434
            log_exc(e)
435
436
            log_debug('No handler for exception')
437
            on_error(e)
438
            remove_all_rows()
439
        # after exception handled, rerun loop with additional constraints
440
441
    if cur != None and row_ct_ref != None and cur.rowcount >= 0:
442
        row_ct_ref[0] += cur.rowcount
443
444
    if is_function: pass # pkeys table already created
445
    elif has_joins:
446
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
447
        log_debug('Getting output table pkeys of existing/inserted rows')
448
        insert_into_pkeys(select_joins, pkeys_cols, distinct=True)
449
    else:
450
        sql.add_row_num(db, insert_out_pkeys) # for joining with input pkeys
451
452
        log_debug('Getting input table pkeys of inserted rows')
453
        sql.run_query_into(db, mk_main_select(input_joins, [in_pkey]),
454
            into=insert_in_pkeys)
455
        sql.add_row_num(db, insert_in_pkeys) # for joining with output pkeys
456
457
        assert sql.table_row_count(db, insert_out_pkeys) == sql.table_row_count(
458
            db, insert_in_pkeys)
459
460
        log_debug('Combining output and input pkeys in inserted order')
461
        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,
462
            {sql.row_num_col: sql_gen.join_same_not_null})]
463
        insert_into_pkeys(pkey_joins, pkeys_names)
464
465
        sql.empty_temp(db, [insert_out_pkeys, insert_in_pkeys])
466
467
    db.log_debug('Adding pkey on pkeys table to enable fast joins', level=2.5)
468
    sql.add_pkey(db, into)
469
470
    log_debug('Setting pkeys of missing rows to '+strings.as_tt(repr(default)))
471
    missing_rows_joins = input_joins+[sql_gen.Join(into,
472
        {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
473
        # must use join_same_not_null or query will take forever
474
    insert_into_pkeys(missing_rows_joins,
475
        [in_pkey_col, sql_gen.NamedCol(out_pkey, default)])
476
477
    assert sql.table_row_count(db, into) == sql.table_row_count(db, in_table)
478
479 3144 aaronmk
    sql.empty_temp(db, insert_in_tables)
480 3077 aaronmk
481
    srcs = []
482
    if is_func: srcs = sql_gen.cols_srcs(in_cols)
483
    return sql_gen.Col(out_pkey, into, srcs)