major-domo-database/app/routers_v3/awards.py

from fastapi import APIRouter, Depends, HTTPException, Query
from typing import List, Optional
import logging
import pydantic

from ..db_engine import db, Award, Team, Player, Manager, model_to_dict, chunked, fn
from ..dependencies import oauth2_scheme, valid_token

router = APIRouter(
    prefix='/api/v3/awards',
    tags=['awards']
)


class AwardModel(pydantic.BaseModel):
    name: str
    season: int
    timing: Optional[str] = "In-Season"
    manager1_id: Optional[int] = None
    manager2_id: Optional[int] = None
    player_id: Optional[int] = None
    team_id: Optional[int] = None
    image: Optional[str] = None


class AwardList(pydantic.BaseModel):
    count: int
    awards: List[AwardModel]


@router.get('')
async def get_awards(
        name: list = Query(default=None), season: Optional[int] = None, timing: Optional[str] = None,
        manager_id: list = Query(default=None), player_id: list = Query(default=None),
        team_id: list = Query(default=None), short_output: Optional[bool] = False):
    all_awards = Award.select()

    if name is not None:
        name_list = [x.lower() for x in name]
        all_awards = all_awards.where(fn.Lower(Award.name) == name_list)
    if season is not None:
        all_awards = all_awards.where(Award.season == season)
    if timing is not None:
        all_awards = all_awards.where(fn.Lower(Award.timing) == timing.lower())
    if manager_id is not None:
        managers = Manager.select().where(Manager.id << manager_id)
        all_awards = all_awards.where(
            (Award.manager1 << managers) | (Award.manager2 << managers)
        )
    if player_id is not None:
        all_awards = all_awards.where(Award.player_id << player_id)
    if team_id is not None:
        all_awards = all_awards.where(Award.team_id << team_id)

    return_awards = {
        'count': all_awards.count(),
        'awards': [model_to_dict(x, recurse=not short_output) for x in all_awards]
    }
    db.close()
    return return_awards


@router.get('/{award_id}')
async def get_one_award(award_id: int, short_output: Optional[bool] = False):
    this_award = Award.get_or_none(Award.id == award_id)
    if this_award is None:
        db.close()
        raise HTTPException(status_code=404, detail=f'Award ID {award_id} not found')

    db.close()
    return model_to_dict(this_award, recurse=not short_output)


@router.patch('/{award_id}')
async def patch_award(
        award_id: int, name: Optional[str] = None, season: Optional[int] = None, timing: Optional[str] = None,
        image: Optional[str] = None, manager1_id: Optional[int] = None, manager2_id: Optional[int] = None,
        player_id: Optional[int] = None, team_id: Optional[int] = None, token: str = Depends(oauth2_scheme)):
    if not valid_token(token):
        logging.warning(f'patch_player - Bad Token: {token}')
        raise HTTPException(status_code=401, detail='Unauthorized')

    this_award = Award.get_or_none(Award.id == award_id)
    if this_award is None:
        db.close()
        raise HTTPException(status_code=404, detail=f'Award ID {award_id} not found')

    if name is not None:
        this_award.name = name
    if season is not None:
        this_award.season = season
    if timing is not None:
        this_award.timing = timing
    if image is not None:
        this_award.image = image
    if manager1_id is not None:
        this_award.manager1_id = manager1_id
    if manager2_id is not None:
        this_award.manager2_id = manager2_id
    if player_id is not None:
        this_award.player_id = player_id
    if team_id is not None:
        this_award.team_id = team_id

    if this_award.save() == 1:
        r_award = model_to_dict(this_award)
        db.close()
        return r_award
    else:
        db.close()
        raise HTTPException(status_code=500, detail=f'Unable to patch award {award_id}')


@router.post('')
async def post_award(award_list: AwardList, token: str = Depends(oauth2_scheme)):
    if not valid_token(token):
        logging.warning(f'patch_player - Bad Token: {token}')
        raise HTTPException(status_code=401, detail='Unauthorized')

    new_awards = []
    for x in award_list.awards:
        if x.manager1_id is not None and Manager.get_or_none(Manager.id == x.manager1_id) is None:
            raise HTTPException(status_code=404, detail=f'Manager ID {x.manager1_id} not found')
        if x.manager2_id is not None and Manager.get_or_none(Manager.id == x.manager2_id) is None:
            raise HTTPException(status_code=404, detail=f'Manager ID {x.manager2_id} not found')
        if x.player_id is not None and Player.get_or_none(Player.id == x.player_id) is None:
            raise HTTPException(status_code=404, detail=f'Player ID {x.player_id} not found')
        if x.team_id is not None and Team.get_or_none(Team.id == x.team_id) is None:
            raise HTTPException(status_code=404, detail=f'Team ID {x.team_id} not found')

        new_awards.append(x.dict())

    with db.atomic():
        for batch in chunked(new_awards, 15):
            Award.insert_many(batch).on_conflict_replace().execute()
    db.close()

    return f'Inserted {len(new_awards)} awards'


@router.delete('/{award_id}')
async def delete_award(award_id: int, token: str = Depends(oauth2_scheme)):
    if not valid_token(token):
        logging.warning(f'patch_player - Bad Token: {token}')
        raise HTTPException(status_code=401, detail='Unauthorized')

    this_award = Award.get_or_none(Award.id == award_id)
    if this_award is None:
        db.close()
        raise HTTPException(status_code=404, detail=f'Award ID {award_id} not found')

    count = this_award.delete_instance()
    db.close()

    if count == 1:
        return f'Award {award_id} has been deleted'
    else:
        raise HTTPException(status_code=500, detail=f'Award {award_id} could not be deleted')