nova-api/api/db/users.py

import os
import yaml
import random
import string
import asyncio

from dotenv import load_dotenv
from motor.motor_asyncio import AsyncIOMotorClient

load_dotenv()

with open('config/credits.yml', encoding='utf8') as f:
    credits_config = yaml.safe_load(f)

## MONGODB Setup

conn = AsyncIOMotorClient(os.getenv('MONGO_URI'))

async def _get_collection(collection_name: str):
    return conn['nova-core'][collection_name]

async def create(discord_id: str='') -> dict:
    """Add a user to the mongodb

    Args:
        discord_id (str): Defaults to ''.

    Returns:
        dict: The user object
    """

    chars = string.ascii_letters + string.digits

    infix = os.getenv('KEYGEN_INFIX')
    suffix = ''.join(random.choices(chars, k=20))
    prefix = ''.join(random.choices(chars, k=20))

    new_api_key = f'nv-{prefix}{infix}{suffix}'

    new_user = {
        'api_key': new_api_key,
        'credits': credits_config['start-credits'],
        'role': '',
        'level': '',
        'status': {
            'active': True,
            'ban_reason': '',
        },
        'auth': {
            'discord': str(discord_id),
            'github': None
        }
    }

    db = await _get_collection('users')
    await db.insert_one(new_user)
    user = await db.find_one({'api_key': new_api_key})	
    return user

async def by_id(user_id: str):
    db = await _get_collection('users')
    return await db.find_one({'_id': user_id})

async def by_discord_id(discord_id: str):
    db = await _get_collection('users')
    return await db.find_one({'auth.discord': str(int(discord_id))})

async def by_api_key(key: str):
    db = await _get_collection('users')
    return await db.find_one({'api_key': key})

async def update_by_id(user_id: str, update):
    db = await _get_collection('users')
    return await db.update_one({'_id': user_id}, update)

async def update_by_filter(obj_filter, update):
    db = await _get_collection('users')
    return await db.update_one(obj_filter, update)

async def delete(user_id: str):
    db = await _get_collection('users')
    await db.delete_one({'_id': user_id})

async def demo():
    user = await create(69420)
    print(user)

if __name__ == '__main__':
    asyncio.run(demo())