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 3077 aaronmk
    conds = set()
314
    distinct_on = sql_gen.ColDict(db, out_table)
315
    def mk_main_select(joins, cols):
316 3106 aaronmk
        conds_ = [(sql_gen.with_table(k, insert_in_table), v) for k, v in conds]
317
        return sql.mk_select(db, joins, cols, conds_, limit=limit_ref[0],
318 3102 aaronmk
            start=0)
319 3077 aaronmk
320
    exc_strs = set()
321
    def log_exc(e):
322
        e_str = exc.str_(e, first_line_only=True)
323
        log_debug('Caught exception: '+e_str)
324
        assert e_str not in exc_strs # avoid infinite loops
325
        exc_strs.add(e_str)
326
327
    def remove_all_rows():
328
        log_debug('Ignoring all rows')
329
        limit_ref[0] = 0 # just create an empty pkeys table
330
331
    def ignore(in_col, value, e):
332 3078 aaronmk
        track_data_error(db, errors_table_, in_col.srcs, value,
333 3077 aaronmk
            e.cause.pgcode, e.cause.pgerror)
334
        log_debug('Ignoring rows with '+strings.as_tt(repr(in_col))+' = '
335
            +strings.as_tt(repr(value)))
336
337
    def remove_rows(in_col, value, e):
338
        ignore(in_col, value, e)
339
        cond = (in_col, sql_gen.CompareCond(value, '!='))
340
        assert cond not in conds # avoid infinite loops
341
        conds.add(cond)
342
343
    def invalid2null(in_col, value, e):
344
        ignore(in_col, value, e)
345
        sql.update(db, in_table, [(in_col, None)],
346
            sql_gen.ColValueCond(in_col, value))
347
348
    def insert_pkeys_table(which):
349
        return sql_gen.Table(sql_gen.concat(in_table.name,
350
            '_insert_'+which+'_pkeys'))
351
    insert_out_pkeys = insert_pkeys_table('out')
352
    insert_in_pkeys = insert_pkeys_table('in')
353
354
    # Do inserts and selects
355
    join_cols = sql_gen.ColDict(db, out_table)
356
    while True:
357
        if limit_ref[0] == 0: # special case
358
            log_debug('Creating an empty pkeys table')
359
            cur = sql.run_query_into(db, sql.mk_select(db, out_table,
360
                [out_pkey], limit=limit_ref[0]), into=insert_out_pkeys)
361
            break # don't do main case
362
363
        has_joins = join_cols != {}
364
365
        log_debug('Trying to insert new rows')
366
367
        # Prepare to insert new rows
368
        insert_args = dict(recover=True, cacheable=False)
369
        if has_joins:
370
            insert_args.update(dict(ignore=True))
371
        else:
372
            insert_args.update(dict(returning=out_pkey, into=insert_out_pkeys))
373 3102 aaronmk
        main_select = mk_main_select([insert_in_table],
374
            [sql_gen.with_table(c, insert_in_table) for c in mapping.values()])
375 3077 aaronmk
376
        def main_insert():
377
            if is_function:
378
                log_debug('Calling function on input rows')
379
                args = dict(((k.name, v) for k, v in mapping.iteritems()))
380
                func_call = sql_gen.NamedCol(out_pkey,
381
                    sql_gen.FunctionCall(out_table, **args))
382
                insert_into_pkeys(input_joins, [in_pkey_col, func_call])
383
                return None
384
            else:
385
                return sql.insert_select(db, out_table, mapping.keys(),
386
                    main_select, **insert_args)
387
388
        try:
389
            cur = sql.with_savepoint(db, main_insert)
390
            break # insert successful
391
        except sql.MissingCastException, e:
392
            log_exc(e)
393
394
            out_col = e.col
395
            type_ = e.type
396
397
            log_debug('Casting '+strings.as_tt(out_col)+' input to '
398
                +strings.as_tt(type_))
399 3111 aaronmk
            in_col = mapping[out_col]
400
            while True:
401
                try:
402
                    mapping[out_col] = cast_temp_col(db, type_, in_col,
403
                        errors_table_)
404
                    break # cast successful
405
                except sql.InvalidValueException, e:
406
                    log_exc(e)
407
408
                    invalid2null(in_col, e.value, e)
409 3077 aaronmk
        except sql.DuplicateKeyException, e:
410
            log_exc(e)
411
412
            old_join_cols = join_cols.copy()
413
            distinct_on.update(util.dict_subset(mapping, e.cols))
414
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
415
            log_debug('Ignoring existing rows, comparing on these columns:\n'
416
                +strings.as_inline_table(join_cols, ustr=col_ustr))
417
            assert join_cols != old_join_cols # avoid infinite loops
418 3102 aaronmk
419
            # Uniquify input table to avoid internal duplicate keys
420
            insert_in_table = sql.distinct_table(db, insert_in_table,
421
                filter(sql_gen.is_table_col, distinct_on.values()))
422 3077 aaronmk
        except sql.NullValueException, e:
423
            log_exc(e)
424
425
            out_col, = e.cols
426
            try: in_col = mapping[out_col]
427
            except KeyError:
428
                log_debug('Missing mapping for NOT NULL column '+out_col)
429
                remove_all_rows()
430
            else: remove_rows(in_col, None, e)
431
        except sql.DatabaseErrors, e:
432
            log_exc(e)
433
434
            log_debug('No handler for exception')
435
            on_error(e)
436
            remove_all_rows()
437
        # after exception handled, rerun loop with additional constraints
438
439
    if cur != None and row_ct_ref != None and cur.rowcount >= 0:
440
        row_ct_ref[0] += cur.rowcount
441
442
    if is_function: pass # pkeys table already created
443
    elif has_joins:
444
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
445
        log_debug('Getting output table pkeys of existing/inserted rows')
446
        insert_into_pkeys(select_joins, pkeys_cols, distinct=True)
447
    else:
448
        sql.add_row_num(db, insert_out_pkeys) # for joining with input pkeys
449
450
        log_debug('Getting input table pkeys of inserted rows')
451
        sql.run_query_into(db, mk_main_select(input_joins, [in_pkey]),
452
            into=insert_in_pkeys)
453
        sql.add_row_num(db, insert_in_pkeys) # for joining with output pkeys
454
455
        assert sql.table_row_count(db, insert_out_pkeys) == sql.table_row_count(
456
            db, insert_in_pkeys)
457
458
        log_debug('Combining output and input pkeys in inserted order')
459
        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,
460
            {sql.row_num_col: sql_gen.join_same_not_null})]
461
        insert_into_pkeys(pkey_joins, pkeys_names)
462
463
        sql.empty_temp(db, [insert_out_pkeys, insert_in_pkeys])
464
465
    db.log_debug('Adding pkey on pkeys table to enable fast joins', level=2.5)
466
    sql.add_pkey(db, into)
467
468
    log_debug('Setting pkeys of missing rows to '+strings.as_tt(repr(default)))
469
    missing_rows_joins = input_joins+[sql_gen.Join(into,
470
        {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
471
        # must use join_same_not_null or query will take forever
472
    insert_into_pkeys(missing_rows_joins,
473
        [in_pkey_col, sql_gen.NamedCol(out_pkey, default)])
474
475
    assert sql.table_row_count(db, into) == sql.table_row_count(db, in_table)
476
477 3102 aaronmk
    sql.empty_temp(db, set([in_table, insert_in_table]))
478 3077 aaronmk
479
    srcs = []
480
    if is_func: srcs = sql_gen.cols_srcs(in_cols)
481
    return sql_gen.Col(out_pkey, into, srcs)