Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac/enh: db session sharing

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2025-12-28 22:00:44 +04:00
parent d4de26bd05
commit 2041ab483e
20 changed files with 600 additions and 562 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/models/messages.py
+43 -42
View File
@@ -3,7 +3,8 @@ import time
import uuid
from typing import Optional
from open_webui.internal.db import Base, get_db
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from open_webui.models.tags import TagModel, Tag, Tags
from open_webui.models.users import Users, User, UserNameResponse
from open_webui.models.channels import Channels, ChannelMember
@@ -137,9 +138,9 @@ class MessageResponse(MessageReplyToResponse):
class MessageTable:
def insert_new_message(
self, form_data: MessageForm, channel_id: str, user_id: str
self, form_data: MessageForm, channel_id: str, user_id: str, db: Optional[Session] = None
) -> Optional[MessageModel]:
with get_db() as db:
with get_db_context(db) as db:
channel_member = Channels.join_channel(channel_id, user_id)
id = str(uuid.uuid4())
@@ -169,22 +170,22 @@ class MessageTable:
db.refresh(result)
return MessageModel.model_validate(result) if result else None
def get_message_by_id(self, id: str) -> Optional[MessageResponse]:
with get_db() as db:
def get_message_by_id(self, id: str, db: Optional[Session] = None) -> Optional[MessageResponse]:
with get_db_context(db) as db:
message = db.get(Message, id)
if not message:
return None
reply_to_message = (
self.get_message_by_id(message.reply_to_id)
self.get_message_by_id(message.reply_to_id, db=db)
if message.reply_to_id
else None
)
reactions = self.get_reactions_by_message_id(id)
thread_replies = self.get_thread_replies_by_message_id(id)
reactions = self.get_reactions_by_message_id(id, db=db)
thread_replies = self.get_thread_replies_by_message_id(id, db=db)
user = Users.get_user_by_id(message.user_id)
user = Users.get_user_by_id(message.user_id, db=db)
return MessageResponse.model_validate(
{
**MessageModel.model_validate(message).model_dump(),
@@ -200,8 +201,8 @@ class MessageTable:
}
)
def get_thread_replies_by_message_id(self, id: str) -> list[MessageReplyToResponse]:
with get_db() as db:
def get_thread_replies_by_message_id(self, id: str, db: Optional[Session] = None) -> list[MessageReplyToResponse]:
with get_db_context(db) as db:
all_messages = (
db.query(Message)
.filter_by(parent_id=id)
@@ -212,7 +213,7 @@ class MessageTable:
messages = []
for message in all_messages:
reply_to_message = (
self.get_message_by_id(message.reply_to_id)
self.get_message_by_id(message.reply_to_id, db=db)
if message.reply_to_id
else None
)
@@ -230,17 +231,17 @@ class MessageTable:
)
return messages
def get_reply_user_ids_by_message_id(self, id: str) -> list[str]:
with get_db() as db:
def get_reply_user_ids_by_message_id(self, id: str, db: Optional[Session] = None) -> list[str]:
with get_db_context(db) as db:
return [
message.user_id
for message in db.query(Message).filter_by(parent_id=id).all()
]
def get_messages_by_channel_id(
self, channel_id: str, skip: int = 0, limit: int = 50
self, channel_id: str, skip: int = 0, limit: int = 50, db: Optional[Session] = None
) -> list[MessageReplyToResponse]:
with get_db() as db:
with get_db_context(db) as db:
all_messages = (
db.query(Message)
.filter_by(channel_id=channel_id, parent_id=None)
@@ -253,7 +254,7 @@ class MessageTable:
messages = []
for message in all_messages:
reply_to_message = (
self.get_message_by_id(message.reply_to_id)
self.get_message_by_id(message.reply_to_id, db=db)
if message.reply_to_id
else None
)
@@ -272,9 +273,9 @@ class MessageTable:
return messages
def get_messages_by_parent_id(
self, channel_id: str, parent_id: str, skip: int = 0, limit: int = 50
self, channel_id: str, parent_id: str, skip: int = 0, limit: int = 50, db: Optional[Session] = None
) -> list[MessageReplyToResponse]:
with get_db() as db:
with get_db_context(db) as db:
message = db.get(Message, parent_id)
if not message:
@@ -296,7 +297,7 @@ class MessageTable:
messages = []
for message in all_messages:
reply_to_message = (
self.get_message_by_id(message.reply_to_id)
self.get_message_by_id(message.reply_to_id, db=db)
if message.reply_to_id
else None
)
@@ -314,8 +315,8 @@ class MessageTable:
)
return messages
def get_last_message_by_channel_id(self, channel_id: str) -> Optional[MessageModel]:
with get_db() as db:
def get_last_message_by_channel_id(self, channel_id: str, db: Optional[Session] = None) -> Optional[MessageModel]:
with get_db_context(db) as db:
message = (
db.query(Message)
.filter_by(channel_id=channel_id)
@@ -325,9 +326,9 @@ class MessageTable:
return MessageModel.model_validate(message) if message else None
def get_pinned_messages_by_channel_id(
self, channel_id: str, skip: int = 0, limit: int = 50
self, channel_id: str, skip: int = 0, limit: int = 50, db: Optional[Session] = None
) -> list[MessageModel]:
with get_db() as db:
with get_db_context(db) as db:
all_messages = (
db.query(Message)
.filter_by(channel_id=channel_id, is_pinned=True)
@@ -339,9 +340,9 @@ class MessageTable:
return [MessageModel.model_validate(message) for message in all_messages]
def update_message_by_id(
self, id: str, form_data: MessageForm
self, id: str, form_data: MessageForm, db: Optional[Session] = None
) -> Optional[MessageModel]:
with get_db() as db:
with get_db_context(db) as db:
message = db.get(Message, id)
message.content = form_data.content
message.data = {
@@ -358,9 +359,9 @@ class MessageTable:
return MessageModel.model_validate(message) if message else None
def update_is_pinned_by_id(
self, id: str, is_pinned: bool, pinned_by: Optional[str] = None
self, id: str, is_pinned: bool, pinned_by: Optional[str] = None, db: Optional[Session] = None
) -> Optional[MessageModel]:
with get_db() as db:
with get_db_context(db) as db:
message = db.get(Message, id)
message.is_pinned = is_pinned
message.pinned_at = int(time.time_ns()) if is_pinned else None
@@ -370,9 +371,9 @@ class MessageTable:
return MessageModel.model_validate(message) if message else None
def get_unread_message_count(
self, channel_id: str, user_id: str, last_read_at: Optional[int] = None
self, channel_id: str, user_id: str, last_read_at: Optional[int] = None, db: Optional[Session] = None
) -> int:
with get_db() as db:
with get_db_context(db) as db:
query = db.query(Message).filter(
Message.channel_id == channel_id,
Message.parent_id == None, # only count top-level messages
@@ -383,9 +384,9 @@ class MessageTable:
return query.count()
def add_reaction_to_message(
self, id: str, user_id: str, name: str
self, id: str, user_id: str, name: str, db: Optional[Session] = None
) -> Optional[MessageReactionModel]:
with get_db() as db:
with get_db_context(db) as db:
# check for existing reaction
existing_reaction = (
db.query(MessageReaction)
@@ -409,8 +410,8 @@ class MessageTable:
db.refresh(result)
return MessageReactionModel.model_validate(result) if result else None
def get_reactions_by_message_id(self, id: str) -> list[Reactions]:
with get_db() as db:
def get_reactions_by_message_id(self, id: str, db: Optional[Session] = None) -> list[Reactions]:
with get_db_context(db) as db:
# JOIN User so all user info is fetched in one query
results = (
db.query(MessageReaction, User)
@@ -440,29 +441,29 @@ class MessageTable:
return [Reactions(**reaction) for reaction in reactions.values()]
def remove_reaction_by_id_and_user_id_and_name(
self, id: str, user_id: str, name: str
self, id: str, user_id: str, name: str, db: Optional[Session] = None
) -> bool:
with get_db() as db:
with get_db_context(db) as db:
db.query(MessageReaction).filter_by(
message_id=id, user_id=user_id, name=name
).delete()
db.commit()
return True
def delete_reactions_by_id(self, id: str) -> bool:
with get_db() as db:
def delete_reactions_by_id(self, id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
db.query(MessageReaction).filter_by(message_id=id).delete()
db.commit()
return True
def delete_replies_by_id(self, id: str) -> bool:
with get_db() as db:
def delete_replies_by_id(self, id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
db.query(Message).filter_by(parent_id=id).delete()
db.commit()
return True
def delete_message_by_id(self, id: str) -> bool:
with get_db() as db:
def delete_message_by_id(self, id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
db.query(Message).filter_by(id=id).delete()
# Delete all reactions to this message