#!/usr/bin/env python3.2

import time

from pgtoolkit import toolbase
from pgtoolkit.highlight import highlight


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
        with self.pgm.cursor('target') as curs:
            rowcount = 1
            while rowcount > 0:
                print('query:', self.args.query)
                curs.execute(self.args.query, [])
                rowcount = curs.rowcount
                print('updated', rowcount)
                curs.connection.commit()
                print('sleep %s seconds' % self.args.sleep)
                time.sleep(self.args.sleep)


tool = BatchUpdateTool()
tool.main()