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,
    PRIVATE_IN_SCHEMA,
    handle_db_errors,
)

logger = logging.getLogger("discord_app")

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]


class DraftPickReturnList(pydantic.BaseModel):
    count: int
    picks: list[DraftPickModel]


@router.get("")
@handle_db_errors
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] = False,
    sort: Optional[str] = None,
    limit: Optional[int] = None,
    player_id: list = Query(default=None),
    player_taken: Optional[bool] = 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 player_id is not None:
        all_picks = all_picks.where(DraftPick.player_id << player_id)
    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 player_taken is not None:
        all_picks = all_picks.where(DraftPick.player.is_null(not player_taken))
    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}")
@handle_db_errors
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}", include_in_schema=PRIVATE_IN_SCHEMA)
@handle_db_errors
async def patch_pick(
    pick_id: int, new_pick: DraftPickModel, token: str = Depends(oauth2_scheme)
):
    if not valid_token(token):
        logger.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("/", include_in_schema=PRIVATE_IN_SCHEMA)
@handle_db_errors
async def post_picks(p_list: DraftPickList, token: str = Depends(oauth2_scheme)):
    if not valid_token(token):
        logger.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_ignore().execute()
    db.close()

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


@router.delete("/{pick_id}", include_in_schema=PRIVATE_IN_SCHEMA)
@handle_db_errors
async def delete_pick(pick_id: int, token: str = Depends(oauth2_scheme)):
    if not valid_token(token):
        logger.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"
        )