refac

backend/open_webui/models/files.py

@@ -135,7 +135,9 @@ class FilesTable:
                 log.exception(f"Error inserting a new file: {e}")
                 return None
 
-    def get_file_by_id(self, id: str, db: Optional[Session] = None) -> Optional[FileModel]:
+    def get_file_by_id(
+        self, id: str, db: Optional[Session] = None
+    ) -> Optional[FileModel]:
         try:
             with get_db_context(db) as db:
                 try:
@@ -146,8 +148,10 @@ class FilesTable:
         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:
+    def get_file_by_id_and_user_id(
+        self, id: str, user_id: str, db: Optional[Session] = None
+    ) -> Optional[FileModel]:
+        with get_db_context(db) as db:
             try:
                 file = db.query(File).filter_by(id=id, user_id=user_id).first()
                 if file:
@@ -157,8 +161,10 @@ class FilesTable:
             except Exception:
                 return None
 
-    def get_file_metadata_by_id(self, id: str) -> Optional[FileMetadataResponse]:
-        with get_db() as db:
+    def get_file_metadata_by_id(
+        self, id: str, db: Optional[Session] = None
+    ) -> Optional[FileMetadataResponse]:
+        with get_db_context(db) as db:
             try:
                 file = db.get(File, id)
                 return FileMetadataResponse(
@@ -175,8 +181,10 @@ class FilesTable:
         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:
-        file = self.get_file_by_id(id)
+    def check_access_by_user_id(
+        self, id, user_id, permission="write", db: Optional[Session] = None
+    ) -> bool:
+        file = self.get_file_by_id(id, db=db)
         if not file:
             return False
         if file.user_id == user_id:
@@ -184,8 +192,10 @@ class FilesTable:
         # Implement additional access control logic here as needed
         return False
 
-    def get_files_by_ids(self, ids: list[str]) -> list[FileModel]:
-        with get_db() as db:
+    def get_files_by_ids(
+        self, ids: list[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)
@@ -194,8 +204,10 @@ class FilesTable:
                 .all()
             ]
 
-    def get_file_metadatas_by_ids(self, ids: list[str]) -> list[FileMetadataResponse]:
-        with get_db() as db:
+    def get_file_metadatas_by_ids(
+        self, ids: list[str], db: Optional[Session] = None
+    ) -> list[FileMetadataResponse]:
+        with get_db_context(db) as db:
             return [
                 FileMetadataResponse(
                     id=file.id,
@@ -212,7 +224,9 @@ class FilesTable:
                 .all()
             ]
 
-    def get_files_by_user_id(self, user_id: str, db: Optional[Session] = None) -> list[FileModel]:
+    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)
@@ -220,9 +234,9 @@ class FilesTable:
             ]
 
     def update_file_by_id(
-        self, id: str, form_data: FileUpdateForm
+        self, id: str, form_data: FileUpdateForm, db: Optional[Session] = None
     ) -> Optional[FileModel]:
-        with get_db() as db:
+        with get_db_context(db) as db:
             try:
                 file = db.query(File).filter_by(id=id).first()
 
@@ -242,8 +256,10 @@ class FilesTable:
                 log.exception(f"Error updating file completely by id: {e}")
                 return None
 
-    def update_file_hash_by_id(self, id: str, hash: str) -> Optional[FileModel]:
-        with get_db() as db:
+    def update_file_hash_by_id(
+        self, id: str, hash: str, db: Optional[Session] = None
+    ) -> Optional[FileModel]:
+        with get_db_context(db) as db:
             try:
                 file = db.query(File).filter_by(id=id).first()
                 file.hash = hash
@@ -254,8 +270,10 @@ class FilesTable:
             except Exception:
                 return None
 
-    def update_file_data_by_id(self, id: str, data: dict) -> Optional[FileModel]:
-        with get_db() as db:
+    def update_file_data_by_id(
+        self, id: str, data: dict, db: Optional[Session] = None
+    ) -> Optional[FileModel]:
+        with get_db_context(db) as db:
             try:
                 file = db.query(File).filter_by(id=id).first()
                 file.data = {**(file.data if file.data else {}), **data}
@@ -266,8 +284,10 @@ class FilesTable:
 
                 return None
 
-    def update_file_metadata_by_id(self, id: str, meta: dict) -> Optional[FileModel]:
-        with get_db() as db:
+    def update_file_metadata_by_id(
+        self, id: str, meta: dict, db: Optional[Session] = None
+    ) -> Optional[FileModel]:
+        with get_db_context(db) as db:
             try:
                 file = db.query(File).filter_by(id=id).first()
                 file.meta = {**(file.meta if file.meta else {}), **meta}
@@ -279,5 +299,25 @@ class FilesTable:
 
                 return False
 
+    def delete_file_by_id(self, id: str, db: Optional[Session] = None) -> bool:
+        with get_db_context(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, db: Optional[Session] = None) -> bool:
+        with get_db_context(db) as db:
+            try:
+                db.query(File).delete()
+                db.commit()
+
+                return True
+            except Exception:
+                return False
+
 
 Files = FilesTable()