refac/enh: db session sharing

backend/open_webui/models/files.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 pydantic import BaseModel, ConfigDict
 from sqlalchemy import BigInteger, Column, String, Text, JSON
 
@@ -108,8 +109,10 @@ class FileListResponse(BaseModel):
 
 
 class FilesTable:
-    def insert_new_file(self, user_id: str, form_data: FileForm) -> Optional[FileModel]:
-        with get_db() as db:
+    def insert_new_file(
+        self, user_id: str, form_data: FileForm, db: Optional[Session] = None
+    ) -> Optional[FileModel]:
+        with get_db_context(db) as db:
             file = FileModel(
                 **{
                     **form_data.model_dump(),
@@ -132,13 +135,16 @@ class FilesTable:
                 log.exception(f"Error inserting a new file: {e}")
                 return None
 
-    def get_file_by_id(self, id: str) -> Optional[FileModel]:
-        with get_db() as db:
-            try:
-                file = db.get(File, id)
-                return FileModel.model_validate(file)
-            except Exception:
-                return None
+    def get_file_by_id(self, id: str, db: Optional[Session] = None) -> Optional[FileModel]:
+        try:
+            with get_db_context(db) as db:
+                try:
+                    file = db.get(File, id)
+                    return FileModel.model_validate(file)
+                except Exception:
+                    return None
+        except Exception:
+            return None
 
     def get_file_by_id_and_user_id(self, id: str, user_id: str) -> Optional[FileModel]:
         with get_db() as db:
@@ -165,8 +171,8 @@ class FilesTable:
             except Exception:
                 return None
 
-    def get_files(self) -> list[FileModel]:
-        with get_db() as db:
+    def get_files(self, db: Optional[Session] = None) -> list[FileModel]:
+        with get_db_context(db) as db:
             return [FileModel.model_validate(file) for file in db.query(File).all()]
 
     def check_access_by_user_id(self, id, user_id, permission="write") -> bool:
@@ -206,8 +212,8 @@ class FilesTable:
                 .all()
             ]
 
-    def get_files_by_user_id(self, user_id: str) -> list[FileModel]:
-        with get_db() as db:
+    def get_files_by_user_id(self, user_id: str, db: Optional[Session] = None) -> list[FileModel]:
+        with get_db_context(db) as db:
             return [
                 FileModel.model_validate(file)
                 for file in db.query(File).filter_by(user_id=user_id).all()
@@ -271,24 +277,6 @@ class FilesTable:
             except Exception:
                 return None
 
-    def delete_file_by_id(self, id: str) -> bool:
-        with get_db() as db:
-            try:
-                db.query(File).filter_by(id=id).delete()
-                db.commit()
-
-                return True
-            except Exception:
-                return False
-
-    def delete_all_files(self) -> bool:
-        with get_db() as db:
-            try:
-                db.query(File).delete()
-                db.commit()
-
-                return True
-            except Exception:
                 return False