refac/enh: db session sharing

backend/open_webui/models/oauth_sessions.py

@@ -8,7 +8,8 @@ import json
 
 from cryptography.fernet import Fernet
 
-from open_webui.internal.db import Base, get_db
+from sqlalchemy.orm import Session
+from open_webui.internal.db import Base, get_db, get_db_context
 from open_webui.env import OAUTH_SESSION_TOKEN_ENCRYPTION_KEY
 
 from pydantic import BaseModel, ConfigDict
@@ -109,10 +110,11 @@ class OAuthSessionTable:
         user_id: str,
         provider: str,
         token: dict,
+        db: Optional[Session] = None,
     ) -> Optional[OAuthSessionModel]:
         """Create a new OAuth session"""
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 current_time = int(time.time())
                 id = str(uuid.uuid4())
 
@@ -141,10 +143,10 @@ class OAuthSessionTable:
             log.error(f"Error creating OAuth session: {e}")
             return None
 
-    def get_session_by_id(self, session_id: str) -> Optional[OAuthSessionModel]:
+    def get_session_by_id(self, session_id: str, db: Optional[Session] = None) -> Optional[OAuthSessionModel]:
         """Get OAuth session by ID"""
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 session = db.query(OAuthSession).filter_by(id=session_id).first()
                 if session:
                     session.token = self._decrypt_token(session.token)
@@ -156,11 +158,11 @@ class OAuthSessionTable:
             return None
 
     def get_session_by_id_and_user_id(
-        self, session_id: str, user_id: str
+        self, session_id: str, user_id: str, db: Optional[Session] = None
     ) -> Optional[OAuthSessionModel]:
         """Get OAuth session by ID and user ID"""
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 session = (
                     db.query(OAuthSession)
                     .filter_by(id=session_id, user_id=user_id)
@@ -176,11 +178,11 @@ class OAuthSessionTable:
             return None
 
     def get_session_by_provider_and_user_id(
-        self, provider: str, user_id: str
+        self, provider: str, user_id: str, db: Optional[Session] = None
     ) -> Optional[OAuthSessionModel]:
         """Get OAuth session by provider and user ID"""
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 session = (
                     db.query(OAuthSession)
                     .filter_by(provider=provider, user_id=user_id)
@@ -195,10 +197,10 @@ class OAuthSessionTable:
             log.error(f"Error getting OAuth session by provider and user ID: {e}")
             return None
 
-    def get_sessions_by_user_id(self, user_id: str) -> List[OAuthSessionModel]:
+    def get_sessions_by_user_id(self, user_id: str, db: Optional[Session] = None) -> List[OAuthSessionModel]:
         """Get all OAuth sessions for a user"""
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 sessions = db.query(OAuthSession).filter_by(user_id=user_id).all()
 
                 results = []
@@ -213,11 +215,11 @@ class OAuthSessionTable:
             return []
 
     def update_session_by_id(
-        self, session_id: str, token: dict
+        self, session_id: str, token: dict, db: Optional[Session] = None
     ) -> Optional[OAuthSessionModel]:
         """Update OAuth session tokens"""
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 current_time = int(time.time())
 
                 db.query(OAuthSession).filter_by(id=session_id).update(
@@ -239,10 +241,10 @@ class OAuthSessionTable:
             log.error(f"Error updating OAuth session tokens: {e}")
             return None
 
-    def delete_session_by_id(self, session_id: str) -> bool:
+    def delete_session_by_id(self, session_id: str, db: Optional[Session] = None) -> bool:
         """Delete an OAuth session"""
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 result = db.query(OAuthSession).filter_by(id=session_id).delete()
                 db.commit()
                 return result > 0
@@ -250,10 +252,10 @@ class OAuthSessionTable:
             log.error(f"Error deleting OAuth session: {e}")
             return False
 
-    def delete_sessions_by_user_id(self, user_id: str) -> bool:
+    def delete_sessions_by_user_id(self, user_id: str, db: Optional[Session] = None) -> bool:
         """Delete all OAuth sessions for a user"""
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 result = db.query(OAuthSession).filter_by(user_id=user_id).delete()
                 db.commit()
                 return True
@@ -261,10 +263,10 @@ class OAuthSessionTable:
             log.error(f"Error deleting OAuth sessions by user ID: {e}")
             return False
 
-    def delete_sessions_by_provider(self, provider: str) -> bool:
+    def delete_sessions_by_provider(self, provider: str, db: Optional[Session] = None) -> bool:
         """Delete all OAuth sessions for a provider"""
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 db.query(OAuthSession).filter_by(provider=provider).delete()
                 db.commit()
                 return True