path: root/scripts/data_maintainer.py

#!/usr/bin/env python

"""Generic script for processing large data sets in small batches.

Reads events from one datasource and commits them into another one,
either one by one or in batches.

Config template::

    [data_maintainer3]
    job_name        = dm_remove_expired_services

    dbread          = dbname=sourcedb_test
    dbwrite         = dbname=destdb port=1234 host=dbhost.com user=guest password=secret
    dbbefore        = dbname=destdb_test
    dbafter         = dbname=destdb_test
    dbcrash         = dbname=destdb_test

    sql_get_pk_list =
        select username
        from user_service
        where expire_date < now();

    # It is a good practice to include same where condition on target side as on read side,
    # to ensure that you are actually changing the same data you think you are,
    # especially when reading from replica database or when processing takes days.
    sql_modify =
        delete from user_service
        where username = %%(username)s
        and expire_date < now();

    # This will be run before executing the sql_get_pk_list query (optional)
    sql_before_run =
        select * from somefunction1(%(job_name)s);

    # This will be run when the DM finishes (optional)
    sql_after_run =
        select * from somefunction2(%(job_name)s);

    # Determines whether the sql_after_run query will be run in case the pk list query returns no rows
    after_zero_rows = 1

    # This will be run if the DM crashes (optional)
    sql_on_crash =
        select * from somefunction3(%(job_name)s);

    # materialize query so that transaction should not be open while processing it
    #with_hold       = 1

    # how many records process to fetch at once and if batch processing is used then
    # also how many records are processed in one commit
    #fetch_count     = 100

    # by default commit after each row (safe when behind plproxy, bouncer or whatever)
    # can be turned off for better performance when connected directly to database
    #autocommit      = 1

    # just for tuning to throttle how much load we let onto write database
    #commit_delay    = 0.0

    # quite often data_maintainer is run from crontab and then loop delay is not needed
    # in case it has to be run as daemon set loop delay in seconds
    #loop_delay      = 1

    logfile         = ~/log/%(job_name)s.log
    pidfile         = ~/pid/%(job_name)s.pid
    use_skylog      = 0
"""

import datetime
import sys
import time

import pkgloader
pkgloader.require('skytools', '3.0')
import skytools


class DataMaintainer (skytools.DBScript):
    __doc__ = __doc__
    loop_delay = -1

    def __init__(self, args):
        super(DataMaintainer, self).__init__("data_maintainer3", args)

        # query for fetching the PK-s of the data set to be maintained
        self.sql_pk = self.cf.get("sql_get_pk_list")

        # query for changing data tuple ( autocommit )
        self.sql_modify = self.cf.get("sql_modify")

        # query to be run before starting the data maintainer,
        # useful for retrieving initialization parameters of the query
        self.sql_before = self.cf.get("sql_before_run", "")

        # query to be run after finishing the data maintainer
        self.sql_after = self.cf.get("sql_after_run", "")

        # whether to run the sql_after query in case of 0 rows
        self.after_zero_rows = self.cf.getint("after_zero_rows", 1)

        # query to be run if the process crashes
        self.sql_crash = self.cf.get("sql_on_crash", "")

        # how many records to fetch at once
        self.fetchcnt = self.cf.getint("fetchcnt", 100)
        self.fetchcnt = self.cf.getint("fetch_count", self.fetchcnt)

        # specifies if non-transactional cursor should be created (0 -> without hold)
        self.withhold = self.cf.getint("with_hold", 1)

        # execution mode (0 -> whole batch is committed / 1 -> autocommit)
        self.autocommit = self.cf.getint("autocommit", 1)

        # delay in seconds after each commit
        self.commit_delay =  self.cf.getfloat("commit_delay", 0.0)

    def work(self):
        self.log.info('Starting..')
        started = lap_time = time.time()
        total_count = 0
        bres = {}

        if self.sql_before:
            bdb = self.get_database("dbbefore", autocommit=1)
            bcur = bdb.cursor()
            bcur.execute(self.sql_before)
            if bcur.statusmessage.startswith('SELECT'):
                res = bcur.dictfetchall()
                assert len(res)==1, "Result of a 'before' query must be 1 row"
                bres = res[0].copy()

        if self.autocommit:
            self.log.info("Autocommit after each modify")
            dbw = self.get_database("dbwrite", autocommit=1)
        else:
            self.log.info("Commit in %i record batches", self.fetchcnt)
            dbw = self.get_database("dbwrite", autocommit=0)
        if self.withhold:
            dbr = self.get_database("dbread", autocommit=1)
            sql = "DECLARE data_maint_cur NO SCROLL CURSOR WITH HOLD FOR %s"
        else:
            dbr = self.get_database("dbread", autocommit=0)
            sql = "DECLARE data_maint_cur NO SCROLL CURSOR FOR %s"
        rcur = dbr.cursor()
        mcur = dbw.cursor()
        rcur.execute(sql % self.sql_pk, bres) # pass results from before_query into sql_pk
        self.log.debug(rcur.query)
        self.log.debug(rcur.statusmessage)

        while True: # loop while fetch returns fetch_count rows
            self.fetch_started = time.time()
            rcur.execute("FETCH FORWARD %s FROM data_maint_cur" % self.fetchcnt)
            self.log.debug(rcur.query)
            self.log.debug(rcur.statusmessage)
            res = rcur.dictfetchall()
            count, lastitem = self.process_batch(res, mcur, bres)
            total_count += count
            if not self.autocommit:
                dbw.commit()
            self.stat_put("duration", time.time() - self.fetch_started)
            self.send_stats()
            if len(res) < self.fetchcnt:
                break
            if not self.looping:
                self.log.info("Exiting on user request")
                break
            if self.commit_delay > 0.0:
                time.sleep(self.commit_delay)
            if time.time() - lap_time > 60.0: # if one minute has passed print running totals
                self.log.info("--- Running count: %s duration: %s ---",
                        total_count,  datetime.timedelta(0, round(time.time() - started)))
                lap_time = time.time()

        rcur.execute("CLOSE data_maint_cur")
        if not self.withhold:
            dbr.rollback()
        self.log.info("--- Total count: %s duration: %s ---",
                total_count,  datetime.timedelta(0, round(time.time() - started)))

        if self.sql_after and (self.after_zero_rows > 0 or total_count > 0):
            adb = self.get_database("dbafter", autocommit=1)
            acur = adb.cursor()

            # FIXME: neither of those can be None?
            if bres != None and lastitem != None:
                bres.update(lastitem)
                lastitem = bres
            if lastitem != None:
                acur.execute(self.sql_after, lastitem)
            else:
                acur.execute(self.sql_after)

    def process_batch(self, res, mcur, bres):
        """ Process events in autocommit mode reading results back and trying to make some sense out of them
        """
        try:
            count = 0
            item = bres
            for i in res:   # for each row in read query result
                item.update(i)
                mcur.execute(self.sql_modify, item)
                self.log.debug(mcur.query)
                if mcur.statusmessage.startswith('SELECT'): # if select was used we can expect some result
                    mres = mcur.dictfetchall()
                    for r in mres:
                        if 'stats' in r: # if specially handled column 'stats' is present
                            for k, v in skytools.db_urldecode(r['stats']).items():
                                self.stat_increase(k, int(v))
                        self.log.debug(r)
                else:
                    self.stat_increase('processed', mcur.rowcount)
                    self.log.debug(mcur.statusmessage)
                if 'cnt' in item:
                    count += item['cnt']
                    self.stat_increase("count", item['cnt'])
                else:
                    count += 1
                    self.stat_increase("count")
                if not self.looping:
                    break
            return count, item
        except: # process has crashed, run sql_crash and re-raise the exception
            if self.sql_crash:
                dbc = self.get_database("dbcrash", autocommit=1)
                ccur = dbc.cursor()
                ccur.execute(self.sql_crash, item)
            raise

if __name__ == '__main__':
    script = DataMaintainer(sys.argv[1:])
    script.start()