bill_man/bill_man.py

#!/usr/bin/env python3

import logging
from datetime import time, datetime

from util import install_except_hook, KSTTZ, UTCTZ


def handle_daemon(name, action, args):
    log = logging.getLogger("main_loop")
    from signal import SIGTERM
    from util.log import wrap

    if action == 'start':
        # load config
        from util import EnteringLoop
        with EnteringLoop(
                name,
                log_dir=args.log_location,
                log_level=args.log_level,
                is_forground=args.foreground,
                pid_path=args.pid_path
        ) as loop:
            from apscheduler.jobstores.memory import MemoryJobStore
            from apscheduler.executors.asyncio import AsyncIOExecutor
            from apscheduler.schedulers.asyncio import AsyncIOScheduler
            from apscheduler.util import undefined
            from signal import SIGTERM, SIGABRT, SIGINT, SIGQUIT
            from handler import Session
            import ujson

            jobstores = {
                'default': MemoryJobStore()
            }
            executors = {
                'default': AsyncIOExecutor(),
            }
            job_defaults = {
                'coalesce': True,
                'max_instances': 2
            }
            scheduler = AsyncIOScheduler(
                jobstores=jobstores,
                executors=executors,
                job_defaults=job_defaults,
                timezone=KSTTZ,
                event_loop=loop)

            with open(args.config) as cfgFile:
                config = ujson.load(cfgFile)
                sess = Session(loop=loop, **config)
                del config
            job = scheduler.add_job(sess.get_resources,
                                    name="charger",
                                    trigger='cron',
                                    day_of_week='mon-fri',
                                    hour=args.standard_time.hour,
                                    minute=args.standard_time.minute,
                                    second=args.standard_time.second,
                                    args=(args.standard_time,),
                                    next_run_time=datetime.utcnow().replace(tzinfo=UTCTZ).astimezone(
                                        KSTTZ) if args.immediately else undefined
                                    )
            scheduler.start()

            if not args.immediately:
                log.info("job(%s) may be start in %s", job.name, job.next_run_time)

            def stopme(*_):
                scheduler.shutdown()
                sess.close()
                log.info("handler closed")
                loop.stop()
                log.info("loop closed")

            loop.add_signal_handler(SIGQUIT, stopme)
            loop.add_signal_handler(SIGTERM, stopme)
            loop.add_signal_handler(SIGINT, stopme)
            loop.add_signal_handler(SIGABRT, stopme)
            loop.run_forever()
        log.info("bye!")
    elif action == 'stop':
        from os import path, kill
        wrap(name, level=args.log_level, stderr=True)
        if not path.exists(args.pid_path):
            log.warning("cannot find pidfile(%s)", args.pid_path)
            return
        with open(args.pid_path, 'r') as pidFile:
            pid = int(pidFile.readline())
            kill(pid, SIGTERM)
            log.warning("pid(%d) sigterm!", pid)
    else:
        raise NotImplementedError()


if __name__ == '__main__':
    import argparse
    from functools import partial

    install_except_hook()

    # 실행 플래그 파싱
    parser = argparse.ArgumentParser(
        prog='bill_man',
        epilog='contact @spi-ca.',
        description='report aws billing .',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--pid_path',
                        help='specify pidPath',
                        default='bill_man.pid')
    parser.add_argument('--log_level',
                        help='set logging level',
                        type=lambda level: logging._nameToLevel[level.upper()],
                        choices=logging._nameToLevel.keys(),
                        default='INFO')

    parser.set_defaults(func=lambda *_: parser.print_help())

    sp = parser.add_subparsers()

    sp_start = sp.add_parser('start', help='Starts %(prog)s daemon',
                             formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    sp_start.add_argument('--config',
                          help='specify config file path',
                          default='bill_man.json')
    sp_start.add_argument('--foreground',
                          help='Don\'t daemonize!',
                          default=False,
                          action='store_true')
    sp_start.add_argument('--immediately',
                          help='run batch now!',
                          default=False,
                          action='store_true')
    sp_start.add_argument('--log_location',
                          help='specify location of logs!',
                          default='log')
    sp_start.add_argument('--standard_time',
                          help='set standard time/HHMMSS',
                          type=lambda ti: time(
                              hour=int(ti[0:2]),
                              minute=int(ti[2:4]),
                              second=int(ti[4:6]),
                              tzinfo=KSTTZ
                          ),
                          default='120000')

    sp_start.set_defaults(func=partial(handle_daemon, parser.prog, 'start'))
    sp_stop = sp.add_parser('stop', help='Stop %(prog)s daemon',
                            formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    sp_stop.set_defaults(func=partial(handle_daemon, parser.prog, 'stop'))

    args = parser.parse_args()
    args.func(args)