Add services/presets.py

Signed-off-by: Vallari Agrawal <[email protected]>

src/teuthology_api/models/presets.py

@@ -1,5 +1,4 @@
 from sqlalchemy import Column, Integer, String, UniqueConstraint
-from sqlalchemy.orm import Session
 from . import Base
 from teuthology_api.schemas.presets import PresetSchema
 
@@ -13,58 +12,3 @@ class Presets(Base):
     cmd = Column(String)
 
     __table_args__ = (UniqueConstraint("username", "name"),)
-
-
-class PresetsDatabaseException(Exception):
-    def __init__(self, message, code):
-        super().__init__(message)
-        self.code = code
-
-
-def create_preset(db: Session, preset):
-    new_preset = Presets(**preset)
-    db.add(new_preset)
-    db.commit()
-    db.refresh(new_preset)
-    return new_preset
-
-
-def get_presets_by_username(db: Session, username: str):
-    db_preset = db.query(Presets).filter(Presets.username == username).all()
-    return db_preset
-
-
-def get_preset_by_username_name(db: Session, username: str, preset_name: str):
-    db_preset = (
-        db.query(Presets)
-        .filter(Presets.username == username, Presets.name == preset_name)
-        .first()
-    )
-    return db_preset
-
-
-def get_preset_id(db: Session, preset_id: int):
-    db_preset = db.query(Presets).filter(Presets.id == preset_id).first()
-    return db_preset
-
-
-def update_preset(db: Session, preset_id: int, update_data):
-    preset_query = db.query(Presets).filter(Presets.id == preset_id)
-    db_preset = preset_query.first()
-    if not db_preset:
-        raise PresetsDatabaseException("Preset does not exist - unable to update.", 404)
-    preset_query.filter(Presets.id == preset_id).update(
-        update_data, synchronize_session=False
-    )
-    db.commit()
-    db.refresh(db_preset)
-    return db_preset
-
-
-def delete_preset(db: Session, preset_id: int):
-    preset_query = db.query(Presets).filter(Presets.id == preset_id)
-    db_preset = preset_query.first()
-    if not db_preset:
-        raise PresetsDatabaseException("Preset does not exist - unable to delete.", 404)
-    preset_query.delete(synchronize_session=False)
-    db.commit()

src/teuthology_api/routes/presets.py

@@ -4,9 +4,8 @@
 
 from teuthology_api.services.helpers import get_token
 from teuthology_api.models import get_db
-from teuthology_api.models.presets import PresetsDatabaseException
-from teuthology_api.models import presets as presets_model
-from teuthology_api.schemas.presets import PresetSchema
+from teuthology_api.services.presets import PresetsDatabaseException, PresetsService
+from teuthology_api.schemas.presets import PresetsSchema
 
 log = logging.getLogger(__name__)
 
@@ -18,23 +17,23 @@
 
 @router.get("/", status_code=200)
 def read_preset(username: str, name: str, db: Session = Depends(get_db)):
-    db_preset = presets_model.get_preset_by_username_name(db, username, name)
+    db_preset = PresetsService(db).get_by_username_and_name(username, name)
     if not db_preset:
         raise HTTPException(status_code=404, detail=f"Preset does not exist.")
     return db_preset
 
 
 @router.get("/list", status_code=200)
 def read_preset(username: str, db: Session = Depends(get_db)):
-    db_presets = presets_model.get_presets_by_username(db, username)
+    db_presets = PresetsService(db).get_by_username(username)
     if not db_presets:
         raise HTTPException(status_code=404, detail=f"User has no presets saved.")
     return db_presets
 
 
 @router.post("/add", status_code=200)
 def add_preset(
-    preset: PresetSchema,
+    preset: PresetsSchema,
     db: Session = Depends(get_db),
     access_token: str = Depends(get_token),
 ):
@@ -44,20 +43,20 @@ def add_preset(
             detail="You need to be logged in",
             headers={"WWW-Authenticate": "Bearer"},
         )
-    db_preset = presets_model.get_preset_by_username_name(
-        db, username=preset.username, preset_name=preset.name
+    db_preset_exists = PresetsService(db).get_by_username_and_name(
+        username=preset.username, preset_name=preset.name
     )
-    if db_preset:
+    if db_preset_exists:
         raise HTTPException(
             status_code=400, detail=f"Preset of this username & name already exists."
         )
-    return presets_model.create_preset(db, preset.model_dump())
+    return PresetsService(db).create(preset.model_dump())
 
 
 @router.put("/edit/{preset_id}", status_code=200)
 def update_preset(
     preset_id: int,
-    updated_data: PresetSchema,
+    updated_data: PresetsSchema,
     db: Session = Depends(get_db),
     access_token: str = Depends(get_token),
 ):
@@ -68,8 +67,8 @@ def update_preset(
             headers={"WWW-Authenticate": "Bearer"},
         )
     try:
-        return presets_model.update_preset(
-            db, preset_id, updated_data.model_dump(exclude_unset=True)
+        return PresetsService(db).update(
+            preset_id, updated_data.model_dump(exclude_unset=True)
         )
     except PresetsDatabaseException as exc:
         raise HTTPException(
@@ -91,7 +90,7 @@ def delete_preset(
             headers={"WWW-Authenticate": "Bearer"},
         )
     try:
-        presets_model.delete_preset(db, preset_id)
+        PresetsService(db).delete(preset_id)
     except PresetsDatabaseException as exc:
         raise HTTPException(
             status_code=exc.code,

src/teuthology_api/schemas/presets.py

@@ -2,10 +2,10 @@
 from typing import Union
 
 
-class PresetSchema(BaseModel):
+class PresetsSchema(BaseModel):
     # pylint: disable=too-few-public-methods
     """
-    Class for Base Args.
+    Class for Presets Schema.
     """
     username: Union[str, None] = Field(default=None)
     name: Union[str, None] = Field(default=None)

src/teuthology_api/services/presets.py

@@ -0,0 +1,60 @@
+from sqlalchemy.orm import Session
+from teuthology_api.models.presets import Presets
+
+
+class PresetsDatabaseException(Exception):
+    def __init__(self, message, code):
+        super().__init__(message)
+        self.code = code
+
+
+class PresetsService:
+    def __init__(self, db: Session) -> None:
+        self.db = db
+
+    def get_by_username(self, username: str):
+        db_preset = self.db.query(Presets).filter(Presets.username == username).all()
+        return db_preset
+
+    def get_by_username_and_name(self, username: str, preset_name: str):
+        db_preset = (
+            self.db.query(Presets)
+            .filter(Presets.username == username, Presets.name == preset_name)
+            .first()
+        )
+        return db_preset
+
+    def get_by_id(self, preset_id: int):
+        db_preset = self.db.query(Presets).filter(Presets.id == preset_id).first()
+        return db_preset
+
+    def create(self, preset: dict) -> Presets:
+        new_preset = Presets(**preset)
+        self.db.add(new_preset)
+        self.db.commit()
+        self.db.refresh(new_preset)
+        return new_preset
+
+    def update(self, preset_id: int, update_data):
+        preset_query = self.db.query(Presets).filter(Presets.id == preset_id)
+        db_preset = preset_query.first()
+        if not db_preset:
+            raise PresetsDatabaseException(
+                "Presets does not exist - unable to update.", 404
+            )
+        preset_query.filter(Presets.id == preset_id).update(
+            update_data, synchronize_session=False
+        )
+        self.db.commit()
+        self.db.refresh(db_preset)
+        return db_preset
+
+    def delete(self, preset_id: int):
+        preset_query = self.db.query(Presets).filter(Presets.id == preset_id)
+        db_preset = preset_query.first()
+        if not db_preset:
+            raise PresetsDatabaseException(
+                "Presets does not exist - unable to delete.", 404
+            )
+        preset_query.delete(synchronize_session=False)
+        self.db.commit()