refac

backend/apps/webui/models/models.py

@@ -5,7 +5,7 @@ from typing import Optional
 from pydantic import BaseModel, ConfigDict
 from sqlalchemy import String, Column, BigInteger, Text
 
-from apps.webui.internal.db import Base, JSONField, Session
+from apps.webui.internal.db import Base, JSONField, get_db
 
 from typing import List, Union, Optional
 from config import SRC_LOG_LEVELS
@@ -126,39 +126,46 @@ class ModelsTable:
             }
         )
         try:
-            result = Model(**model.model_dump())
-            Session.add(result)
-            Session.commit()
-            Session.refresh(result)
 
-            if result:
-                return ModelModel.model_validate(result)
-            else:
-                return None
+            with get_db() as db:
+
+                result = Model(**model.model_dump())
+                db.add(result)
+                db.commit()
+                db.refresh(result)
+
+                if result:
+                    return ModelModel.model_validate(result)
+                else:
+                    return None
         except Exception as e:
             print(e)
             return None
 
     def get_all_models(self) -> List[ModelModel]:
-        return [
-            ModelModel.model_validate(model) for model in Session.query(Model).all()
-        ]
+        with get_db() as db:
+
+            return [ModelModel.model_validate(model) for model in db.query(Model).all()]
 
     def get_model_by_id(self, id: str) -> Optional[ModelModel]:
         try:
-            model = Session.get(Model, id)
-            return ModelModel.model_validate(model)
+            with get_db() as db:
+
+                model = db.get(Model, id)
+                return ModelModel.model_validate(model)
         except:
             return None
 
     def update_model_by_id(self, id: str, model: ModelForm) -> Optional[ModelModel]:
         try:
-            # update only the fields that are present in the model
-            model = Session.query(Model).get(id)
-            model.update(**model.model_dump())
-            Session.commit()
-            Session.refresh(model)
-            return ModelModel.model_validate(model)
+            with get_db() as db:
+
+                # update only the fields that are present in the model
+                model = db.query(Model).get(id)
+                model.update(**model.model_dump())
+                db.commit()
+                db.refresh(model)
+                return ModelModel.model_validate(model)
         except Exception as e:
             print(e)
 
@@ -166,8 +173,10 @@ class ModelsTable:
 
     def delete_model_by_id(self, id: str) -> bool:
         try:
-            Session.query(Model).filter_by(id=id).delete()
-            return True
+            with get_db() as db:
+
+                db.query(Model).filter_by(id=id).delete()
+                return True
         except:
             return False