refac/enh: db session sharing

backend/open_webui/models/tools.py

@@ -2,7 +2,8 @@ import logging
 import time
 from typing import Optional
 
-from open_webui.internal.db import Base, JSONField, get_db
+from sqlalchemy.orm import Session
+from open_webui.internal.db import Base, JSONField, get_db, get_db_context
 from open_webui.models.users import Users, UserResponse
 from open_webui.models.groups import Groups
 
@@ -110,9 +111,9 @@ class ToolValves(BaseModel):
 
 class ToolsTable:
     def insert_new_tool(
-        self, user_id: str, form_data: ToolForm, specs: list[dict]
+        self, user_id: str, form_data: ToolForm, specs: list[dict], db: Optional[Session] = None
     ) -> Optional[ToolModel]:
-        with get_db() as db:
+        with get_db_context(db) as db:
             tool = ToolModel(
                 **{
                     **form_data.model_dump(),
@@ -136,21 +137,21 @@ class ToolsTable:
                 log.exception(f"Error creating a new tool: {e}")
                 return None
 
-    def get_tool_by_id(self, id: str) -> Optional[ToolModel]:
+    def get_tool_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ToolModel]:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 tool = db.get(Tool, id)
                 return ToolModel.model_validate(tool)
         except Exception:
             return None
 
-    def get_tools(self) -> list[ToolUserModel]:
-        with get_db() as db:
+    def get_tools(self, db: Optional[Session] = None) -> list[ToolUserModel]:
+        with get_db_context(db) as db:
             all_tools = db.query(Tool).order_by(Tool.updated_at.desc()).all()
 
             user_ids = list(set(tool.user_id for tool in all_tools))
 
-            users = Users.get_users_by_user_ids(user_ids) if user_ids else []
+            users = Users.get_users_by_user_ids(user_ids, db=db) if user_ids else []
             users_dict = {user.id: user for user in users}
 
             tools = []
@@ -167,10 +168,10 @@ class ToolsTable:
             return tools
 
     def get_tools_by_user_id(
-        self, user_id: str, permission: str = "write"
+        self, user_id: str, permission: str = "write", db: Optional[Session] = None
     ) -> list[ToolUserModel]:
-        tools = self.get_tools()
-        user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user_id)}
+        tools = self.get_tools(db=db)
+        user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user_id, db=db)}
 
         return [
             tool
@@ -179,31 +180,31 @@ class ToolsTable:
             or has_access(user_id, permission, tool.access_control, user_group_ids)
         ]
 
-    def get_tool_valves_by_id(self, id: str) -> Optional[dict]:
+    def get_tool_valves_by_id(self, id: str, db: Optional[Session] = None) -> Optional[dict]:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 tool = db.get(Tool, id)
                 return tool.valves if tool.valves else {}
         except Exception as e:
             log.exception(f"Error getting tool valves by id {id}")
             return None
 
-    def update_tool_valves_by_id(self, id: str, valves: dict) -> Optional[ToolValves]:
+    def update_tool_valves_by_id(self, id: str, valves: dict, db: Optional[Session] = None) -> Optional[ToolValves]:
         try:
-            with get_db() as db:
+            with get_db_context(db) 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)
+                return self.get_tool_by_id(id, db=db)
         except Exception:
             return None
 
     def get_user_valves_by_id_and_user_id(
-        self, id: str, user_id: str
+        self, id: str, user_id: str, db: Optional[Session] = None
     ) -> Optional[dict]:
         try:
-            user = Users.get_user_by_id(user_id)
+            user = Users.get_user_by_id(user_id, db=db)
             user_settings = user.settings.model_dump() if user.settings else {}
 
             # Check if user has "tools" and "valves" settings
@@ -220,10 +221,10 @@ class ToolsTable:
             return None
 
     def update_user_valves_by_id_and_user_id(
-        self, id: str, user_id: str, valves: dict
+        self, id: str, user_id: str, valves: dict, db: Optional[Session] = None
     ) -> Optional[dict]:
         try:
-            user = Users.get_user_by_id(user_id)
+            user = Users.get_user_by_id(user_id, db=db)
             user_settings = user.settings.model_dump() if user.settings else {}
 
             # Check if user has "tools" and "valves" settings
@@ -235,7 +236,7 @@ class ToolsTable:
             user_settings["tools"]["valves"][id] = valves
 
             # Update the user settings in the database
-            Users.update_user_by_id(user_id, {"settings": user_settings})
+            Users.update_user_by_id(user_id, {"settings": user_settings}, db=db)
 
             return user_settings["tools"]["valves"][id]
         except Exception as e:
@@ -244,9 +245,9 @@ class ToolsTable:
             )
             return None
 
-    def update_tool_by_id(self, id: str, updated: dict) -> Optional[ToolModel]:
+    def update_tool_by_id(self, id: str, updated: dict, db: Optional[Session] = None) -> Optional[ToolModel]:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 db.query(Tool).filter_by(id=id).update(
                     {**updated, "updated_at": int(time.time())}
                 )
@@ -258,9 +259,9 @@ class ToolsTable:
         except Exception:
             return None
 
-    def delete_tool_by_id(self, id: str) -> bool:
+    def delete_tool_by_id(self, id: str, db: Optional[Session] = None) -> bool:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 db.query(Tool).filter_by(id=id).delete()
                 db.commit()