"""PgQ cooperative consumer for Python.
"""

from pgq.consumer import Consumer

__all__ = ['CoopConsumer']

class CoopConsumer(Consumer):
    """Cooperative Consumer base class.

    There will be one dbscript process per subconsumer.

    Config params::
        ## pgq.CoopConsumer

        # name for subconsumer
        subconsumer_name =

        # pgsql interval when to consider parallel subconsumers dead,
        # and take over their unfinished batch
        #subconsumer_timeout = 1 hour
    """

    def __init__(self, service_name, db_name, args):
        """Initialize new subconsumer.

        @param service_name: service_name for DBScript
        @param db_name: name of database for get_database()
        @param args: cmdline args for DBScript
        """

        Consumer.__init__(self, service_name, db_name, args)

        self.subconsumer_name = self.cf.get("subconsumer_name")
        self.subconsumer_timeout = self.cf.get("subconsumer_timeout", "")

    def register_consumer(self):
        """Registration for subconsumer."""

        self.log.info("Registering consumer on source queue")
        db = self.get_database(self.db_name)
        cx = db.cursor()
        cx.execute("select pgq_coop.register_subconsumer(%s, %s, %s)",
                [self.queue_name, self.consumer_name, self.subconsumer_name])
        res = cx.fetchone()[0]
        db.commit()

        return res

    def unregister_consumer(self):
        """Unregistration for subconsumer."""

        self.log.info("Unregistering consumer from source queue")
        db = self.get_database(self.db_name)
        cx = db.cursor()
        cx.execute("select pgq_coop.unregister_subconsumer(%s, %s, %s, 0)",
                    [self.queue_name, self.consumer_name, self.subconsumer_name])
        db.commit()


    def _load_next_batch(self, curs):
        """Allocate next batch. (internal)"""

        if self.subconsumer_timeout:
            q = "select pgq_coop.next_batch(%s, %s, %s, %s)"
            curs.execute(q, [self.queue_name, self.consumer_name, self.subconsumer_name, self.subconsumer_timeout])
        else:
            q = "select pgq_coop.next_batch(%s, %s, %s)"
            curs.execute(q, [self.queue_name, self.consumer_name, self.subconsumer_name])
        return curs.fetchone()[0]

    def _finish_batch(self, curs, batch_id, list):
        """Finish batch. (internal)"""

        self._flush_retry(curs, batch_id, list)
        curs.execute("select pgq_coop.finish_batch(%s)", [batch_id])