feat(sqlalchemy): remove session reference from router

This commit is contained in:
Jonathan Rohde
2024-06-21 14:58:57 +02:00
parent df09d0830a
commit bee835cb65
34 changed files with 1231 additions and 1211 deletions

View File

@@ -5,7 +5,7 @@ import logging
from sqlalchemy import String, Column, BigInteger
from sqlalchemy.orm import Session
from apps.webui.internal.db import Base, JSONField
from apps.webui.internal.db import Base, JSONField, get_session
from apps.webui.models.users import Users
import json
@@ -82,7 +82,7 @@ class ToolValves(BaseModel):
class ToolsTable:
def insert_new_tool(
self, db: Session, user_id: str, form_data: ToolForm, specs: List[dict]
self, user_id: str, form_data: ToolForm, specs: List[dict]
) -> Optional[ToolModel]:
tool = ToolModel(
**{
@@ -95,46 +95,48 @@ class ToolsTable:
)
try:
result = Tool(**tool.dict())
db.add(result)
db.commit()
db.refresh(result)
if result:
return ToolModel.model_validate(result)
else:
return None
with get_session() as db:
result = Tool(**tool.model_dump())
db.add(result)
db.commit()
db.refresh(result)
if result:
return ToolModel.model_validate(result)
else:
return None
except Exception as e:
print(f"Error creating tool: {e}")
return None
def get_tool_by_id(self, db: Session, id: str) -> Optional[ToolModel]:
def get_tool_by_id(self, id: str) -> Optional[ToolModel]:
try:
tool = db.get(Tool, id)
return ToolModel.model_validate(tool)
with get_session() as db:
tool = db.get(Tool, id)
return ToolModel.model_validate(tool)
except:
return None
def get_tools(self, db: Session) -> List[ToolModel]:
return [ToolModel.model_validate(tool) for tool in db.query(Tool).all()]
def get_tools(self) -> List[ToolModel]:
with get_session() as db:
return [ToolModel.model_validate(tool) for tool in db.query(Tool).all()]
def get_tool_valves_by_id(self, id: str) -> Optional[dict]:
try:
tool = Tool.get(Tool.id == id)
return tool.valves if tool.valves else {}
with get_session() as db:
tool = db.get(Tool, id)
return tool.valves if tool.valves else {}
except Exception as e:
print(f"An error occurred: {e}")
return None
def update_tool_valves_by_id(self, id: str, valves: dict) -> Optional[ToolValves]:
try:
query = Tool.update(
**{"valves": valves},
updated_at=int(time.time()),
).where(Tool.id == id)
query.execute()
tool = Tool.get(Tool.id == id)
return ToolValves(**model_to_dict(tool))
with get_session() as db:
db.query(Tool).filter_by(id=id).update(
{"valves": valves, "updated_at": int(time.time())}
)
db.commit()
return self.get_tool_by_id(id)
except:
return None
@@ -172,8 +174,7 @@ class ToolsTable:
user_settings["tools"]["valves"][id] = valves
# Update the user settings in the database
query = Users.update_user_by_id(user_id, {"settings": user_settings})
query.execute()
Users.update_user_by_id(user_id, {"settings": user_settings})
return user_settings["tools"]["valves"][id]
except Exception as e:
@@ -182,16 +183,19 @@ class ToolsTable:
def update_tool_by_id(self, id: str, updated: dict) -> Optional[ToolModel]:
try:
db.query(Tool).filter_by(id=id).update(
{**updated, "updated_at": int(time.time())}
)
return self.get_tool_by_id(db, id)
with get_session() as db:
db.query(Tool).filter_by(id=id).update(
{**updated, "updated_at": int(time.time())}
)
db.commit()
return self.get_tool_by_id(id)
except:
return None
def delete_tool_by_id(self, db: Session, id: str) -> bool:
def delete_tool_by_id(self, id: str) -> bool:
try:
db.query(Tool).filter_by(id=id).delete()
with get_session() as db:
db.query(Tool).filter_by(id=id).delete()
return True
except:
return False