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

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

from ..db_engine import db, DraftPick, Team, model_to_dict, chunked
from ..dependencies import oauth2_scheme, valid_token

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


class DraftPickModel(pydantic.BaseModel):
    overall: Optional[int] = None
    round: int
    origowner_id: int
    owner_id: Optional[int] = None
    season: int
    player_id: Optional[int] = None


class DraftPickList(pydantic.BaseModel):
    picks: List[DraftPickModel]


@router.get('')
async def get_picks(
        season: int, owner_team_abbrev: list = Query(default=None), orig_team_abbrev: list = Query(default=None),
        owner_team_id: list = Query(default=None), orig_team_id: list = Query(default=None),
        pick_round_start: Optional[int] = None, pick_round_end: Optional[int] = None, traded: Optional[bool] = None,
        overall: Optional[int] = None, overall_start: Optional[int] = None, overall_end: Optional[int] = None,
        short_output: Optional[bool] = True, sort: Optional[str] = None, limit: Optional[int] = None):
    all_picks = DraftPick.select().where(DraftPick.season == season)

    if owner_team_abbrev is not None:
        team_list = []
        for x in owner_team_abbrev:
            team_list.append(Team.get_season(x, season))
        all_picks = all_picks.where(
            (DraftPick.owner << team_list) | (DraftPick.owner << team_list)
        )

    if orig_team_abbrev is not None:
        team_list = []
        for x in orig_team_abbrev:
            team_list.append(Team.get_season(x, season))
        all_picks = all_picks.where(
            (DraftPick.origowner << team_list) | (DraftPick.origowner << team_list)
        )

    if owner_team_id is not None:
        all_picks = all_picks.where(
            (DraftPick.owner_id << owner_team_id) | (DraftPick.owner_id << owner_team_id)
        )

    if orig_team_id is not None:
        all_picks = all_picks.where(
            (DraftPick.origowner_id << orig_team_id) | (DraftPick.origowner_id << orig_team_id)
        )

    if pick_round_start is not None and pick_round_end is not None and pick_round_end < pick_round_start:
        raise HTTPException(status_code=400, detail=f'pick_round_end must be greater than or equal to pick_round_start')

    if pick_round_start is not None:
        all_picks = all_picks.where(DraftPick.round >= pick_round_start)
    if pick_round_end is not None:
        all_picks = all_picks.where(DraftPick.round <= pick_round_end)
    if traded is not None:
        all_picks = all_picks.where(DraftPick.origowner != DraftPick.owner)
    if overall is not None:
        all_picks = all_picks.where(DraftPick.overall == overall)
    if overall_start is not None:
        all_picks = all_picks.where(DraftPick.overall >= overall_start)
    if overall_end is not None:
        all_picks = all_picks.where(DraftPick.overall <= overall_end)
    if limit is not None:
        all_picks = all_picks.limit(limit)

    if sort is not None:
        if sort == 'order-asc':
            all_picks = all_picks.order_by(DraftPick.overall)
        elif sort == 'order-desc':
            all_picks = all_picks.order_by(-DraftPick.overall)

    return_picks = {'count': all_picks.count(), 'picks': []}
    for line in all_picks:
        return_picks['picks'].append(model_to_dict(line, recurse=not short_output))

    db.close()
    return return_picks


@router.get('/{pick_id}')
async def get_one_pick(pick_id: int, short_output: Optional[bool] = False):
    this_pick = DraftPick.get_or_none(DraftPick.id == pick_id)
    if this_pick is not None:
        r_pick = model_to_dict(this_pick, recurse=not short_output)
    else:
        raise HTTPException(status_code=404, detail=f'Pick ID {pick_id} not found')
    db.close()
    return r_pick


@router.patch('/{pick_id}')
async def patch_pick(pick_id: int, new_pick: DraftPickModel, token: str = Depends(oauth2_scheme)):
    if not valid_token(token):
        logging.warning(f'patch_pick - Bad Token: {token}')
        raise HTTPException(status_code=401, detail='Unauthorized')

    if DraftPick.get_or_none(DraftPick.id == pick_id) is None:
        raise HTTPException(status_code=404, detail=f'Pick ID {pick_id} not found')

    DraftPick.update(**new_pick.dict()).where(DraftPick.id == pick_id).execute()
    r_pick = model_to_dict(DraftPick.get_by_id(pick_id))
    db.close()
    return r_pick


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

    new_picks = []
    for pick in p_list.picks:
        dupe = DraftPick.get_or_none(DraftPick.season == pick.season, DraftPick.overall == pick.overall)
        if dupe:
            db.close()
            raise HTTPException(
                status_code=500,
                detail=f'Pick # {pick.overall} already exists for season {pick.season}'
            )

        new_picks.append(pick.dict())

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

    return f'Inserted {len(new_picks)} picks'


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

    this_pick = DraftPick.get_or_none(DraftPick.id == pick_id)
    if this_pick is None:
        raise HTTPException(status_code=404, detail=f'Pick ID {pick_id} not found')

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

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