From f0fa121a4851146619eae5fa5d23b732f623ce45 Mon Sep 17 00:00:00 2001 From: Devansh Singh Date: Tue, 13 Aug 2024 18:52:00 +0530 Subject: [PATCH] Add PresetCreate schema, update Presets cmd column type Added a schema for /add endpoint request body for validating cmd attribute. Updated the cmd column datatype from string to JSON for the Presets model. Signed-off-by: Devansh Singh --- .../versions/db0c0dc793f4_add_presets_table.py | 2 +- src/teuthology_api/models/presets.py | 7 +++++-- src/teuthology_api/routes/presets.py | 17 +++++++++++------ src/teuthology_api/schemas/preset.py | 12 ++++++++++++ 4 files changed, 29 insertions(+), 9 deletions(-) create mode 100644 src/teuthology_api/schemas/preset.py diff --git a/migrations/versions/db0c0dc793f4_add_presets_table.py b/migrations/versions/db0c0dc793f4_add_presets_table.py index df60e7d..26c73e9 100644 --- a/migrations/versions/db0c0dc793f4_add_presets_table.py +++ b/migrations/versions/db0c0dc793f4_add_presets_table.py @@ -28,7 +28,7 @@ def upgrade() -> None: sa.Column("username", sqlmodel.sql.sqltypes.AutoString(), nullable=False), sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False), sa.Column("suite", sqlmodel.sql.sqltypes.AutoString(), nullable=False), - sa.Column("cmd", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("cmd", sa.JSON(), nullable=False), sa.PrimaryKeyConstraint("id"), sa.UniqueConstraint("username", "name"), ) diff --git a/src/teuthology_api/models/presets.py b/src/teuthology_api/models/presets.py index cdd7f17..7a3aebb 100644 --- a/src/teuthology_api/models/presets.py +++ b/src/teuthology_api/models/presets.py @@ -1,11 +1,14 @@ -from sqlmodel import Field, SQLModel, UniqueConstraint +from pydantic import ConfigDict +from sqlmodel import Column, Field, JSON, SQLModel, UniqueConstraint class Presets(SQLModel, table=True): + model_config = ConfigDict(arbitrary_types_allowed=True) + id: int = Field(primary_key=True) username: str = Field(index=True) name: str suite: str - cmd: str + cmd: JSON = Field(sa_column=Column(JSON)) __table_args__ = (UniqueConstraint("username", "name"),) diff --git a/src/teuthology_api/routes/presets.py b/src/teuthology_api/routes/presets.py index 0087aad..0f51cd3 100644 --- a/src/teuthology_api/routes/presets.py +++ b/src/teuthology_api/routes/presets.py @@ -1,10 +1,11 @@ import logging -from fastapi import status, APIRouter, HTTPException, Depends +from fastapi import status, APIRouter, Depends, HTTPException, Request from sqlmodel import Session -from teuthology_api.services.helpers import get_token from teuthology_api.models import get_db, Presets +from teuthology_api.schemas.preset import PresetCreate +from teuthology_api.services.helpers import get_token, get_username from teuthology_api.services.presets import PresetsDatabaseException, PresetsService logger = logging.getLogger(__name__) @@ -36,7 +37,8 @@ def read_all_presets(username: str, db: Session = Depends(get_db)): @router.post("/add", status_code=status.HTTP_201_CREATED) def add_preset( - preset: Presets, + request: Request, + preset: PresetCreate, replace: bool = False, db: Session = Depends(get_db), access_token: str = Depends(get_token), @@ -48,7 +50,8 @@ def add_preset( headers={"WWW-Authenticate": "Bearer"}, ) - db_presets = PresetsService(db).get_by_username(preset.username) + username = get_username(request) + db_presets = PresetsService(db).get_by_username(username) if len(db_presets) == 10: if not replace: raise HTTPException( @@ -60,14 +63,16 @@ def add_preset( PresetsService(db).delete(db_presets[0].id) db_preset_exists = PresetsService(db).get_by_username_and_name( - preset.username, preset.name + username, preset.name ) if db_preset_exists: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Preset with name {preset.name} exists", ) - return PresetsService(db).create(preset) + + db_preset = Presets(**preset.model_dump(), username=username) + return PresetsService(db).create(db_preset) @router.put("/edit/{preset_id}", status_code=status.HTTP_200_OK) diff --git a/src/teuthology_api/schemas/preset.py b/src/teuthology_api/schemas/preset.py new file mode 100644 index 0000000..958f4e6 --- /dev/null +++ b/src/teuthology_api/schemas/preset.py @@ -0,0 +1,12 @@ +from pydantic import BaseModel + +from teuthology_api.schemas.suite import SuiteArgs + + +class PresetCreate(BaseModel): + name: str + suite: str + cmd: SuiteArgs + + def model_post_init(self, __context): + self.cmd = self.cmd.model_dump()