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
        cur = sql.insert(db, table, row, pkey_, recover=True)
171
        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
        return sql.value(sql.select(db, table, [pkey_], row, recover=True))
178
179
def get(db, table, row, pkey, row_ct_ref=None, create=False):
180
    '''Recovers from errors'''
181
    try:
182
        return sql.value(sql.select(db, table, [pkey], row, limit=1,
183
            recover=True))
184
    except StopIteration:
185
        if not create: raise
186
        return put(db, table, row, pkey, row_ct_ref) # insert new row
187
188
def is_func_result(col):
189
    return col.table.name.find('(') >= 0 and col.name == 'result'
190
191
def into_table_name(out_table, in_tables0, mapping, is_func):
192
    def in_col_str(in_col):
193
        in_col = sql_gen.remove_col_rename(in_col)
194
        if isinstance(in_col, sql_gen.Col):
195
            table = in_col.table
196
            if table == in_tables0:
197
                in_col = sql_gen.to_name_only_col(in_col)
198
            elif is_func_result(in_col): in_col = table # omit col name
199
        return str(in_col)
200
201
    str_ = str(out_table)
202
    if is_func:
203
        str_ += '('
204
205
        try: value_in_col = mapping['value']
206
        except KeyError:
207
            str_ += ', '.join((str(k)+'='+in_col_str(v)
208
                for k, v in mapping.iteritems()))
209
        else: str_ += in_col_str(value_in_col)
210
211
        str_ += ')'
212
    else:
213
        out_col = 'rank'
214
        try: in_col = mapping[out_col]
215
        except KeyError: str_ += '_pkeys'
216
        else: # has a rank column, so hierarchical
217
            str_ += '['+str(out_col)+'='+in_col_str(in_col)+']'
218
    return str_
219
220
def put_table(db, out_table, in_tables, mapping, row_ct_ref=None, into=None,
221
    default=None, is_func=False, on_error=exc.raise_):
222
    '''Recovers from errors.
223
    Only works under PostgreSQL (uses INSERT RETURNING).
224
    IMPORTANT: Must be run at the *beginning* of a transaction.
225
    @param in_tables The main input table to select from, followed by a list of
226
        tables to join with it using the main input table's pkey
227
    @param mapping dict(out_table_col=in_table_col, ...)
228
        * out_table_col: str (*not* sql_gen.Col)
229
        * in_table_col: sql_gen.Col|literal-value
230
    @param into The table to contain the output and input pkeys.
231
        Defaults to `out_table.name+'_pkeys'`.
232
    @param default The *output* column to use as the pkey for missing rows.
233
        If this output column does not exist in the mapping, uses None.
234
    @param is_func Whether out_table is the name of a SQL function, not a table
235
    @return sql_gen.Col Where the output pkeys are made available
236
    '''
237
    out_table = sql_gen.as_Table(out_table)
238
239
    def log_debug(msg): db.log_debug(msg, level=1.5)
240
    def col_ustr(str_):
241
        return strings.repr_no_u(sql_gen.remove_col_rename(str_))
242
243
    log_debug('********** New iteration **********')
244
    log_debug('Inserting these input columns into '+strings.as_tt(
245
        out_table.to_str(db))+':\n'+strings.as_table(mapping, ustr=col_ustr))
246
247
    is_function = sql.function_exists(db, out_table)
248
249
    if is_function: out_pkey = 'result'
250
    else: out_pkey = sql.pkey(db, out_table, recover=True)
251
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
252
253
    if mapping == {}: # need at least one column for INSERT SELECT
254
        mapping = {out_pkey: None} # ColDict will replace with default value
255
256
    # Create input joins from list of input tables
257
    in_tables_ = in_tables[:] # don't modify input!
258
    in_tables0 = in_tables_.pop(0) # first table is separate
259 3078 aaronmk
    errors_table_ = errors_table(db, in_tables0)
260 3077 aaronmk
    in_pkey = sql.pkey(db, in_tables0, recover=True)
261
    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
262
    input_joins = [in_tables0]+[sql_gen.Join(v,
263
        {in_pkey: sql_gen.join_same_not_null}) for v in in_tables_]
264
265
    if into == None:
266
        into = into_table_name(out_table, in_tables0, mapping, is_func)
267
    into = sql_gen.as_Table(into)
268
269
    # Set column sources
270
    in_cols = filter(sql_gen.is_table_col, mapping.values())
271
    for col in in_cols:
272
        if col.table == in_tables0: col.set_srcs(sql_gen.src_self)
273
274
    log_debug('Joining together input tables into temp table')
275
    # Place in new table for speed and so don't modify input if values edited
276
    in_table = sql_gen.Table('in')
277
    mapping = dicts.join(mapping, sql.flatten(db, in_table, input_joins,
278
        in_cols, preserve=[in_pkey_col], start=0))
279
    input_joins = [in_table]
280
    db.log_debug('Temp table: '+strings.as_tt(in_table.to_str(db)), level=2)
281
282
    mapping = sql_gen.ColDict(db, out_table, mapping)
283
        # after applying dicts.join() because that returns a plain dict
284
285
    # Resolve default value column
286
    if default != None:
287
        try: default = mapping[default]
288
        except KeyError:
289
            db.log_debug('Default value column '
290
                +strings.as_tt(strings.repr_no_u(default))
291
                +' does not exist in mapping, falling back to None', level=2.1)
292
            default = None
293
294
    pkeys_names = [in_pkey, out_pkey]
295
    pkeys_cols = [in_pkey_col, out_pkey_col]
296
297
    pkeys_table_exists_ref = [False]
298
    def insert_into_pkeys(joins, cols, distinct=False):
299
        kw_args = {}
300
        if distinct: kw_args.update(dict(distinct_on=[in_pkey_col]))
301
        query = sql.mk_select(db, joins, cols, order_by=None, start=0,
302
            **kw_args)
303
304
        if pkeys_table_exists_ref[0]:
305
            sql.insert_select(db, into, pkeys_names, query)
306
        else:
307
            sql.run_query_into(db, query, into=into)
308
            pkeys_table_exists_ref[0] = True
309
310
    limit_ref = [None]
311 3106 aaronmk
    insert_in_table = in_table
312 3077 aaronmk
    conds = set()
313
    distinct_on = sql_gen.ColDict(db, out_table)
314
    def mk_main_select(joins, cols):
315 3106 aaronmk
        conds_ = [(sql_gen.with_table(k, insert_in_table), v) for k, v in conds]
316
        return sql.mk_select(db, joins, cols, conds_, limit=limit_ref[0],
317 3102 aaronmk
            start=0)
318 3077 aaronmk
319
    exc_strs = set()
320
    def log_exc(e):
321
        e_str = exc.str_(e, first_line_only=True)
322
        log_debug('Caught exception: '+e_str)
323
        assert e_str not in exc_strs # avoid infinite loops
324
        exc_strs.add(e_str)
325
326
    def remove_all_rows():
327
        log_debug('Ignoring all rows')
328
        limit_ref[0] = 0 # just create an empty pkeys table
329
330
    def ignore(in_col, value, e):
331 3078 aaronmk
        track_data_error(db, errors_table_, in_col.srcs, value,
332 3077 aaronmk
            e.cause.pgcode, e.cause.pgerror)
333
        log_debug('Ignoring rows with '+strings.as_tt(repr(in_col))+' = '
334
            +strings.as_tt(repr(value)))
335
336
    def remove_rows(in_col, value, e):
337
        ignore(in_col, value, e)
338
        cond = (in_col, sql_gen.CompareCond(value, '!='))
339
        assert cond not in conds # avoid infinite loops
340
        conds.add(cond)
341
342
    def invalid2null(in_col, value, e):
343
        ignore(in_col, value, e)
344
        sql.update(db, in_table, [(in_col, None)],
345
            sql_gen.ColValueCond(in_col, value))
346
347
    def insert_pkeys_table(which):
348
        return sql_gen.Table(sql_gen.concat(in_table.name,
349
            '_insert_'+which+'_pkeys'))
350
    insert_out_pkeys = insert_pkeys_table('out')
351
    insert_in_pkeys = insert_pkeys_table('in')
352
353
    # Do inserts and selects
354
    join_cols = sql_gen.ColDict(db, out_table)
355
    while True:
356
        if limit_ref[0] == 0: # special case
357
            log_debug('Creating an empty pkeys table')
358
            cur = sql.run_query_into(db, sql.mk_select(db, out_table,
359
                [out_pkey], limit=limit_ref[0]), into=insert_out_pkeys)
360
            break # don't do main case
361
362
        has_joins = join_cols != {}
363
364
        log_debug('Trying to insert new rows')
365
366
        # Prepare to insert new rows
367
        insert_args = dict(recover=True, cacheable=False)
368
        if has_joins:
369
            insert_args.update(dict(ignore=True))
370
        else:
371
            insert_args.update(dict(returning=out_pkey, into=insert_out_pkeys))
372 3102 aaronmk
        main_select = mk_main_select([insert_in_table],
373
            [sql_gen.with_table(c, insert_in_table) for c in mapping.values()])
374 3077 aaronmk
375
        def main_insert():
376
            if is_function:
377
                log_debug('Calling function on input rows')
378
                args = dict(((k.name, v) for k, v in mapping.iteritems()))
379
                func_call = sql_gen.NamedCol(out_pkey,
380
                    sql_gen.FunctionCall(out_table, **args))
381
                insert_into_pkeys(input_joins, [in_pkey_col, func_call])
382
                return None
383
            else:
384
                return sql.insert_select(db, out_table, mapping.keys(),
385
                    main_select, **insert_args)
386
387
        try:
388
            cur = sql.with_savepoint(db, main_insert)
389
            break # insert successful
390
        except sql.MissingCastException, e:
391
            log_exc(e)
392
393
            out_col = e.col
394
            type_ = e.type
395
396
            log_debug('Casting '+strings.as_tt(out_col)+' input to '
397
                +strings.as_tt(type_))
398 3111 aaronmk
            in_col = mapping[out_col]
399
            while True:
400
                try:
401
                    mapping[out_col] = cast_temp_col(db, type_, in_col,
402
                        errors_table_)
403
                    break # cast successful
404
                except sql.InvalidValueException, e:
405
                    log_exc(e)
406
407
                    invalid2null(in_col, e.value, e)
408 3077 aaronmk
        except sql.DuplicateKeyException, e:
409
            log_exc(e)
410
411
            old_join_cols = join_cols.copy()
412
            distinct_on.update(util.dict_subset(mapping, e.cols))
413
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
414
            log_debug('Ignoring existing rows, comparing on these columns:\n'
415
                +strings.as_inline_table(join_cols, ustr=col_ustr))
416
            assert join_cols != old_join_cols # avoid infinite loops
417 3102 aaronmk
418
            # Uniquify input table to avoid internal duplicate keys
419
            insert_in_table = sql.distinct_table(db, insert_in_table,
420
                filter(sql_gen.is_table_col, distinct_on.values()))
421 3077 aaronmk
        except sql.NullValueException, e:
422
            log_exc(e)
423
424
            out_col, = e.cols
425
            try: in_col = mapping[out_col]
426
            except KeyError:
427
                log_debug('Missing mapping for NOT NULL column '+out_col)
428
                remove_all_rows()
429
            else: remove_rows(in_col, None, e)
430
        except sql.DatabaseErrors, e:
431
            log_exc(e)
432
433
            log_debug('No handler for exception')
434
            on_error(e)
435
            remove_all_rows()
436
        # after exception handled, rerun loop with additional constraints
437
438
    if cur != None and row_ct_ref != None and cur.rowcount >= 0:
439
        row_ct_ref[0] += cur.rowcount
440
441
    if is_function: pass # pkeys table already created
442
    elif has_joins:
443
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
444
        log_debug('Getting output table pkeys of existing/inserted rows')
445
        insert_into_pkeys(select_joins, pkeys_cols, distinct=True)
446
    else:
447
        sql.add_row_num(db, insert_out_pkeys) # for joining with input pkeys
448
449
        log_debug('Getting input table pkeys of inserted rows')
450
        sql.run_query_into(db, mk_main_select(input_joins, [in_pkey]),
451
            into=insert_in_pkeys)
452
        sql.add_row_num(db, insert_in_pkeys) # for joining with output pkeys
453
454
        assert sql.table_row_count(db, insert_out_pkeys) == sql.table_row_count(
455
            db, insert_in_pkeys)
456
457
        log_debug('Combining output and input pkeys in inserted order')
458
        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,
459
            {sql.row_num_col: sql_gen.join_same_not_null})]
460
        insert_into_pkeys(pkey_joins, pkeys_names)
461
462
        sql.empty_temp(db, [insert_out_pkeys, insert_in_pkeys])
463
464
    db.log_debug('Adding pkey on pkeys table to enable fast joins', level=2.5)
465
    sql.add_pkey(db, into)
466
467
    log_debug('Setting pkeys of missing rows to '+strings.as_tt(repr(default)))
468
    missing_rows_joins = input_joins+[sql_gen.Join(into,
469
        {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
470
        # must use join_same_not_null or query will take forever
471
    insert_into_pkeys(missing_rows_joins,
472
        [in_pkey_col, sql_gen.NamedCol(out_pkey, default)])
473
474
    assert sql.table_row_count(db, into) == sql.table_row_count(db, in_table)
475
476 3102 aaronmk
    sql.empty_temp(db, set([in_table, insert_in_table]))
477 3077 aaronmk
478
    srcs = []
479
    if is_func: srcs = sql_gen.cols_srcs(in_cols)
480
    return sql_gen.Col(out_pkey, into, srcs)