diff --git a/main.py b/main.py index 9cfefdf..d634452 100644 --- a/main.py +++ b/main.py @@ -7,10 +7,11 @@ import json from db.constants import TIMELINE_ORDER_ISSUED, TIMELINE_ORDER_NOT_ISSUED from scheduling import OrderScheduler from generate import generate_order -from orders import order_issue, order_check -from db.queries import orders_pool_by_id, timeline_event_put +from orders import order_issue, order_check, punishment_issue +from db.queries import orders_pool_by_id, order_status_by_id from telegram.telegram import handle_commands from telegram.commands import commands +from util import make_session logger = logging.getLogger(__name__) @@ -19,6 +20,12 @@ async def do_order_issue(orders_pool_id): await order_issue(orders_pool) +async def do_punishment_issue(order_status_id): + order_status = order_status_by_id(order_status_id) + + async with make_session() as session: + await punishment_issue(session, order_status) + if __name__=='__main__': logging.basicConfig( format="%(asctime)s %(module)s [%(levelname)-4.4s] %(message)s", @@ -40,6 +47,9 @@ if __name__=='__main__': parser_issue = subparsers.add_parser('check', help='Check on the status of an order') parser_issue.add_argument('order_status_id') + parser_punish = subparsers.add_parser('punish', help="Punish an order") + parser_punish.add_argument('order_status_id') + args = parser.parse_args() if args.command == 'generate': @@ -56,6 +66,11 @@ if __name__=='__main__': loop.run_until_complete( order_check(args.order_status_id) ) + elif args.command == 'punish': + loop = asyncio.new_event_loop() + loop.run_until_complete( + do_punishment_issue(args.order_status_id) + ) else: loop = asyncio.new_event_loop() s = OrderScheduler(loop)