Simple restart manager

fabfile.py

# Example fabric task

import elasticsearch
from restartmanager import RestartManager

@task
def do_the_restart():
    # fab qa:es-master do_the_restart
    r = RestartManager("/path/to/db")
    r.schedule(*env.hosts)
    es = elasticsearch.Elasticsearch("elasticsearch.com:9200")
    while True:
        host = r.current
        if not host:
            host = r.start()
            # sudo("service elasticsearch restart")
        while True:
            health = es.cluster.health(wait_for_status="green", request_timeout=10)
            if health["status"] == "green" and health["unassigned_shards"] == 0:
                r.finish()
                break

restartmanager.py

import sqlite3
import datetime
import calendar
import collections

Restart = collections.namedtuple("Restart",
    ["name", "position", "started_at", "finished_at"])

def to_restart(row):
    r = dict(row)
    if r["started_at"] is not None:
        r["started_at"] = datetime.datetime.utcfromtimestamp(r["started_at"])
    if r["finished_at"] is not None:
        r["finished_at"] = datetime.datetime.utcfromtimestamp(r["finished_at"])
    return r

def returns(klass, converter=None):
    if converter is None:
        converter = dict

    def decorator(fn):
        def inner(*args, **kwargs):
            res = fn(*args, **kwargs)
            if isinstance(res, sqlite3.Row):
                return klass(**converter(res))
            if not isinstance(res, (list, tuple)):
                return res
            return [klass(**converter(r)) for r in res]
        return inner
    return decorator

class RestartManager(object):
    def __init__(self, path=":memory:"):
        self.db = sqlite3.connect(path)
        self.db.row_factory = sqlite3.Row
        self.init_db()

    def execute(self, sql, args=None, commit=False, many=False):
        cur = self.db.cursor()
        fn = getattr(cur, "executemany" if many else "execute")
        if args:
            res = fn(sql, args)
        else:
            res = fn(sql)
        if commit:
            self.db.commit()
        return res

    def reset(self):
        sql = """
        delete from reservations
        """
        self.execute(sql, commit=True)

    def init_db(self):
        sql = """
        create table if not exists restarts(
            name text primary key not null,
            position integer not null,
            started_at timestamp,
            finished_at timestamp
        );
        """
        self.execute(sql, commit=True)

    @property
    def current_timestamp(self):
        now = datetime.datetime.utcnow()
        return calendar.timegm(now.timetuple())

    @property
    @returns(Restart, to_restart)
    def current(self):
        sql = """
        select * from restarts
        where started_at is not null
          and finished_at is null
        order by position
        limit 1
        """
        res = self.execute(sql)
        return res.fetchone()

    @property
    @returns(Restart, to_restart)
    def next(self):
        sql = """
        select * from restarts
        where started_at is null
          and finished_at is null
        order by position
        limit 1
        """
        res = self.execute(sql)
        return res.fetchone()

    @property
    @returns(Restart, to_restart)
    def prev(self):
        sql = """
        select * from restarts
        where started_at is not null
          and finished_at is not null
        order by position desc
        limit 1
        """
        res = self.execute(sql)
        return res.fetchone()

    @property
    @returns(Restart, to_restart)
    def remaining(self):
        sql = """
        select * from restarts
        where finished_at is null
        order by position
        """
        res = self.execute(sql)
        return res.fetchall()

    @property
    @returns(Restart, to_restart)
    def all(self):
        sql = """
        select * from restarts
        order by position
        """
        return self.execute(sql).fetchall()

    def schedule(self, *names):
        sql = """
        insert or ignore into restarts(name, position)
        values(:name, (select coalesce(max(position), 0) from restarts) + 1)
        """
        res = self.execute(sql, [{"name": n} for n in names], True, True)
        return res

    def start(self, name=None):
        if name is None and self.current:
            name = self.current.name
        elif name is None and self.next:
            name = self.next.name
        if name is None:
            raise RuntimeError("must supply a name to start")
        sql = """
        update restarts
        set started_at = coalesce(started_at, :now)
        where name = :name
        """
        args = {"name": name, "now": self.current_timestamp}
        self.execute(sql, args, commit=True)
        return self.current

    def finish(self, name=None):
        if name is None and self.current:
            name = self.current.name
        if name is None:
            raise RuntimeError("must supply a name to finish")
        sql = """
        update restarts
        set finished_at = coalesce(finished_at, :now)
        where name = :name
        """
        args = {"name": name, "now": self.current_timestamp}
        return self.execute(sql, args, commit=True)

def main():
    resman = RestartManager()

if __name__ == "__main__":
  main()