diff --git a/manage/cmd/users.py b/manage/cmd/users.py index efd3cf4..48717a0 100644 --- a/manage/cmd/users.py +++ b/manage/cmd/users.py @@ -94,6 +94,42 @@ async def adduser(ctx, args): print(f'\tdiscrim: {user["discriminator"]}') +async def set_max_concurrency(ctx, args): + """Update the `max_concurrency` for a bot. + This can only be set for bot accounts! + """ + + if int(args.max_concurrency) < 1: + return print("max_concurrency must be >0") + + bot = await ctx.db.fetchval( + """ + select bot + from users + where id = $1 + """, + int(args.user_id), + ) + + if bot == None: + return print("user not found") + + if bot == False: + return print("user must be a bot") + + await ctx.db.execute( + """ + update users + set max_concurrency = $1 + where id = $2 + """, + int(args.max_concurrency), + int(args.user_id), + ) + + print(f"OK: set max_concurrency={args.max_concurrency} for {args.user_id}") + + async def set_flag(ctx, args): """Setting a 'staff' flag gives the user access to the Admin API. Beware of that. @@ -198,6 +234,17 @@ def setup(subparser): setup_test_parser.set_defaults(func=adduser) + set_max_concurrency_parser = subparser.add_parser( + "set_max_concurrency", + help="set `max_concurrency` for a user", + description=set_max_concurrency.__doc__, + ) + set_max_concurrency_parser.add_argument("user_id") + set_max_concurrency_parser.add_argument( + "max_concurrency", help="the `max_concurrency` value to set" + ) + set_max_concurrency_parser.set_defaults(func=set_max_concurrency) + setflag_parser = subparser.add_parser( "setflag", help="set a flag for a user", description=set_flag.__doc__ )