Add PresetArgs 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 <[email protected]>

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"),
     )

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"),)

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 PresetArgs
+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: PresetArgs,
     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,20 +63,22 @@ 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)
 def update_preset(
     preset_id: int,
-    updated_preset: Presets,
+    updated_preset: PresetArgs,
     db: Session = Depends(get_db),
     access_token: str = Depends(get_token),
 ):

src/teuthology_api/schemas/preset.py

@@ -0,0 +1,12 @@
+from pydantic import BaseModel
+
+from teuthology_api.schemas.suite import SuiteArgs
+
+
+class PresetArgs(BaseModel):
+    name: str
+    suite: str
+    cmd: SuiteArgs
+
+    def model_post_init(self, __context):
+        self.cmd = self.cmd.model_dump()

tests/test_preset.py

@@ -1,17 +1,22 @@
+from unittest.mock import patch
 from urllib.parse import urlencode
 
 from fastapi.testclient import TestClient
 from sqlmodel import Session
 
 preset_payload = {
-    "username": "user1",
     "name": "test",
     "suite": "teuthology:no-ceph",
-    "cmd": '{"owner": "user1"}',
+    "cmd": {
+        "--suite": "teuthology:no-ceph",
+        "--owner": "user1",
+    },
 }
 
 
-def test_create_preset(session: Session, client: TestClient):
+@patch("teuthology_api.routes.presets.get_username")
+def test_create_preset(m_get_username, session: Session, client: TestClient):
+    m_get_username.return_value = "user1"
     response = client.post("/presets/add", json=preset_payload)
     assert response.status_code == 201