Răsfoiți Sursa

refac/fix: proper notes db operations

Timothy Jaeryang Baek 1 săptămână în urmă
părinte
comite
da661756fa
2 a modificat fișierele cu 37 adăugiri și 31 ștergeri
  1. 33 29
      backend/open_webui/models/notes.py
  2. 4 2
      backend/open_webui/routers/notes.py

+ 33 - 29
backend/open_webui/models/notes.py

@@ -128,7 +128,7 @@ class NoteTable:
             notes = query.all()
             return [NoteModel.model_validate(note) for note in notes]
 
-    def get_notes_by_access(
+    def get_notes_by_permission(
         self,
         user_id: str,
         permission: str = "write",
@@ -137,40 +137,44 @@ class NoteTable:
     ) -> list[NoteModel]:
         with get_db() as db:
             user_groups = Groups.get_groups_by_member_id(user_id)
-            user_group_ids = {group_id for group_id in user_groups}
-
-            query = db.query(Note)
-
-            access_conditions = [Note.user_id == user_id]
+            user_group_ids = {group.id for group in user_groups}
+
+            # Order newest-first. We stream to keep memory usage low.
+            query = (
+                db.query(Note)
+                .order_by(Note.updated_at.desc())
+                .execution_options(stream_results=True)
+                .yield_per(256)
+            )
 
-            if user_group_ids:
-                access_conditions.append(
-                    and_(
-                        Note.access_control.isnot(None),
-                        Note.access_control != '{}',
-                        Note.access_control != 'null'
+            results: list[NoteModel] = []
+            n_skipped = 0
+
+            for note in query:
+                # Fast-pass #1: owner
+                if note.user_id == user_id:
+                    permitted = True
+                # Fast-pass #2: public/open
+                elif note.access_control is None:
+                    permitted = True
+                else:
+                    permitted = has_access(
+                        user_id, permission, note.access_control, user_group_ids
                     )
-                )
 
-            query = query.filter(or_(*access_conditions))
+                if not permitted:
+                    continue
 
-            query = query.order_by(Note.updated_at.desc())
-
-            if skip is not None:
-                query = query.offset(skip)
-            if limit is not None:
-                query = query.limit(limit)
-
-            notes = query.all()
-            note_models = [NoteModel.model_validate(note) for note in notes]
+                # Apply skip AFTER permission filtering so it counts only accessible notes
+                if skip and n_skipped < skip:
+                    n_skipped += 1
+                    continue
 
-            filtered_notes = []
-            for note in note_models:
-                if (note.user_id == user_id or
-                    has_access(user_id, permission, note.access_control, user_group_ids)):
-                    filtered_notes.append(note)
+                results.append(NoteModel.model_validate(note))
+                if limit is not None and len(results) >= limit:
+                    break
 
-            return filtered_notes
+            return results
 
     def get_note_by_id(self, id: str) -> Optional[NoteModel]:
         with get_db() as db:

+ 4 - 2
backend/open_webui/routers/notes.py

@@ -48,7 +48,7 @@ async def get_notes(request: Request, user=Depends(get_verified_user)):
                 "user": UserResponse(**Users.get_user_by_id(note.user_id).model_dump()),
             }
         )
-        for note in Notes.get_notes_by_access(user.id, "write")
+        for note in Notes.get_notes_by_permission(user.id, "write")
     ]
 
     return notes
@@ -81,7 +81,9 @@ async def get_note_list(
 
     notes = [
         NoteTitleIdResponse(**note.model_dump())
-        for note in Notes.get_notes_by_access(user.id, "write", skip=skip, limit=limit)
+        for note in Notes.get_notes_by_permission(
+            user.id, "write", skip=skip, limit=limit
+        )
     ]
 
     return notes