#!/usr/bin/env python3.2

import time

from pgtoolkit import toolbase

from pgtoolkit.highlight import highlight

from pgtoolkit.pgmanager import OperationalError

class BatchUpdateTool(toolbase.SimpleTool):

    def __init__(self):

        toolbase.SimpleTool.__init__(self, name='batchupdate', desc='Run an UPDATE query until it does not update any rows.')

        self.parser.add_argument('--query', dest='query', type=str, help='Query to run. It should be constructed using subquery with LIMIT.')

        self.parser.add_argument('--sleep', dest='sleep', default=0, type=int, help='Sleep SLEEP seconds between subsequent queries.')

        self.init()

    def main(self):

        # connect DB

        rowcount = 1

        while rowcount > 0:

            print('query:', self.args.query)

            with self.pgm.cursor('target') as curs:

                try:

                    curs.execute(self.args.query, [])

                    rowcount = curs.rowcount

                    print('updated', rowcount)

                    curs.connection.commit()

                except (OperationalError, SystemError) as e:

                    print('Error:', str(e))

            print('sleep %s seconds' % self.args.sleep)

            time.sleep(self.args.sleep)

tool = BatchUpdateTool()

tool.main()