1 #!/usr/bin/env python3

     2 

     3 from pgtoolkit import toolbase

     4 

     5 

     6 class BatchCopyTool(toolbase.SrcDstTablesTool):

     7     def __init__(self):

     8         toolbase.SrcDstTablesTool.__init__(self, name='batchcopy', desc='Copy data from one table to another.')

     9 

    10         self.parser.add_argument('--table-name', type=str, help='Table to be copied.')

    11         self.parser.add_argument('--src-filter', type=str, help='WHERE condition for source query.')

    12         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}')

    13         #TODO: duplicates=rollback|ignore|update

    14 

    15         self.init()

    16 

    17     def main(self):

    18         # read list of IDs from file

    19         ids = '<no IDs read>'

    20         if self.args.file_with_ids:

    21             with open(self.args.file_with_ids, 'r') as f:

    22                 ids = ','.join(ln.rstrip() for ln in f.readlines())

    23 

    24         # read source data

    25         with self.pgm.cursor('src') as src_curs:

    26             condition = self.args.src_filter.format(ids=ids) or 'true'

    27             src_curs.execute('SELECT * FROM {} WHERE {}'.format(self.args.table_name, condition))

    28             #TODO:  ORDER BY id OFFSET 0 LIMIT 100

    29             data = src_curs.fetchall_dict()

    30             src_curs.connection.commit()

    31 

    32         with self.pgm.cursor('dst') as dst_curs:

    33             for row in data:

    34                 keys = ', '.join(row.keys())

    35                 values_mask = ', '.join(['%s'] * len(row))

    36                 query = 'INSERT INTO {} ({}) VALUES ({})'.format(self.args.table_name, keys, values_mask)

    37                 dst_curs.execute(query, list(row.values()))

    38             dst_curs.connection.commit()

    39 

    40         self.log.info('Copied {} rows.' % len(rows))

    41 

    42 

    43 tool = BatchCopyTool()

    44 tool.main()