diff --git a/manage/cmd/migration/command.py b/manage/cmd/migration/command.py index a1e87c6..4f80b64 100644 --- a/manage/cmd/migration/command.py +++ b/manage/cmd/migration/command.py @@ -1,4 +1,5 @@ import inspect +import os from pathlib import Path from dataclasses import dataclass from collections import namedtuple @@ -22,14 +23,13 @@ class MigrationContext: @property def latest(self): """Return the latest migration ID.""" - return max(self.scripts.keys()) - + return 0 if len(self.scripts) == 0 else max(self.scripts.keys()) def make_migration_ctx() -> MigrationContext: """Create the MigrationContext instance.""" # taken from https://stackoverflow.com/a/6628348 script_path = inspect.stack()[0][1] - script_folder = '/'.join(script_path.split('/')[:-1]) + script_folder = os.sep.join(script_path.split(os.sep)[:-1]) script_folder = Path(script_folder) migration_folder = script_folder / 'scripts' @@ -40,7 +40,7 @@ def make_migration_ctx() -> MigrationContext: mig_path_str = str(mig_path) # extract migration script id and name - mig_filename = mig_path_str.split('/')[-1].split('.')[0] + mig_filename = mig_path_str.split(os.sep)[-1].split('.')[0] name_fragments = mig_filename.split('_') mig_id = int(name_fragments[0]) diff --git a/manage/cmd/users.py b/manage/cmd/users.py index 92701d5..8887fa5 100644 --- a/manage/cmd/users.py +++ b/manage/cmd/users.py @@ -12,14 +12,6 @@ async def find_user(username, discrim, ctx): WHERE username = $1 AND discriminator = $2 """, username, discrim) -async def get_password_hash(id, ctx): - return await ctx.db.fetchval(""" - SELECT password_hash - FROM users - WHERE id = $1 - """, id) - - async def set_user_staff(user_id, ctx): """Give a single user staff status.""" old_flags = await ctx.db.fetchval(""" @@ -62,15 +54,19 @@ async def make_staff(ctx, args): await set_user_staff(uid, ctx) print('OK: set staff') -async def generate_token(ctx, args): - """Generate a token for specified user.""" - uid = await find_user(args.username, args.discrim, ctx) +async def generate_bot_token(ctx, args): + """Generate a token for specified bot.""" - if not uid: - return print('user not found') - - password_hash = await get_password_hash(uid, ctx) - print(make_token(uid, password_hash)) + password_hash = await ctx.db.fetchval(""" + SELECT password_hash + FROM users + WHERE id = $1 AND bot = 'true' + """, int(args.user_id)) + + if not password_hash: + return print('cannot find a bot with specified id') + + print(make_token(args.user_id, password_hash)) def setup(subparser): @@ -105,15 +101,12 @@ def setup(subparser): token_parser = subparser.add_parser( 'generate_token', - help='generate a token for specified user', - description=generate_token.__doc__ + help='generate a token for specified bot', + description=generate_bot_token.__doc__ ) token_parser.add_argument( - 'username' - ) - token_parser.add_argument( - 'discrim', help='the discriminator of the user' + 'user_id' ) - token_parser.set_defaults(func=generate_token) + token_parser.set_defaults(func=generate_bot_token)