batchcopy.py
author Radek Brich <brich.radek@ifortuna.cz>
Tue, 06 May 2014 18:34:38 +0200
changeset 99 245646538743
parent 98 024299702087
permissions -rwxr-xr-x
Update runquery tool: Add --one-query-per-line parameter.
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
97
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
     1
#!/usr/bin/env python3
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
     2
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
     3
from pgtoolkit import toolbase
98
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
     4
from pgtoolkit.pgmanager import IntegrityError
97
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
     5
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
     6
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
     7
class BatchCopyTool(toolbase.SrcDstTablesTool):
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
     8
    def __init__(self):
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
     9
        toolbase.SrcDstTablesTool.__init__(self, name='batchcopy', desc='Copy data from one table to another.')
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    10
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    11
        self.parser.add_argument('--table-name', type=str, help='Table to be copied.')
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    12
        self.parser.add_argument('--src-filter', type=str, help='WHERE condition for source query.')
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    13
        self.parser.add_argument('--file-with-ids', type=str, help='Read source IDs from file (each ID on new line). Use these in --src-filter as {ids}')
98
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    14
        self.parser.add_argument('--dst-exists', choices=['rollback', 'ignore', 'update'], default='rollback', help='What to do when destination record already exists.')
97
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    15
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    16
        self.init()
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    17
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    18
    def main(self):
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    19
        # read list of IDs from file
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    20
        ids = '<no IDs read>'
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    21
        if self.args.file_with_ids:
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    22
            with open(self.args.file_with_ids, 'r') as f:
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    23
                ids = ','.join(ln.rstrip() for ln in f.readlines())
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    24
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    25
        # read source data
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    26
        with self.pgm.cursor('src') as src_curs:
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    27
            condition = self.args.src_filter.format(ids=ids) or 'true'
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    28
            src_curs.execute('SELECT * FROM {} WHERE {}'.format(self.args.table_name, condition))
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    29
            #TODO:  ORDER BY id OFFSET 0 LIMIT 100
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    30
            data = src_curs.fetchall_dict()
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    31
            src_curs.connection.commit()
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    32
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    33
        with self.pgm.cursor('dst') as dst_curs:
98
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    34
            copied = 0
97
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    35
            for row in data:
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    36
                keys = ', '.join(row.keys())
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    37
                values_mask = ', '.join(['%s'] * len(row))
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    38
                query = 'INSERT INTO {} ({}) VALUES ({})'.format(self.args.table_name, keys, values_mask)
98
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    39
                try:
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    40
                    dst_curs.execute('SAVEPOINT the_query;')
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    41
                    dst_curs.execute(query, list(row.values()))
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    42
                    dst_curs.execute('RELEASE SAVEPOINT the_query;')
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    43
                    copied += 1
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    44
                except IntegrityError:
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    45
                    if self.args.dst_exists == 'rollback':
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    46
                        dst_curs.connection.rollback()
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    47
                        break
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    48
                    elif self.args.dst_exists == 'ignore':
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    49
                        dst_curs.execute('ROLLBACK TO SAVEPOINT the_query;')
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    50
                    elif self.args.dst_exists == 'update':
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    51
                        raise NotImplementedError()
97
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    52
            dst_curs.connection.commit()
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    53
98
024299702087 Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents: 97
diff changeset
    54
        self.log.info('Copied %s rows.', copied)
97
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    55
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    56
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    57
tool = BatchCopyTool()
a4af93e72e2b Add batchcopy tool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
diff changeset
    58
tool.main()