refac/enh: db session sharing

backend/open_webui/models/auths.py

@@ -2,7 +2,8 @@ import logging
 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.users import UserModel, UserProfileImageResponse, Users
 from pydantic import BaseModel
 from sqlalchemy import Boolean, Column, String, Text
@@ -87,8 +88,9 @@ class AuthsTable:
         profile_image_url: str = "/user.png",
         role: str = "pending",
         oauth: Optional[dict] = None,
+        db: Optional[Session] = None,
     ) -> Optional[UserModel]:
-        with get_db() as db:
+        with get_db_context(db) as db:
             log.info("insert_new_auth")
 
             id = str(uuid.uuid4())
@@ -100,7 +102,7 @@ class AuthsTable:
             db.add(result)
 
             user = Users.insert_new_user(
-                id, name, email, profile_image_url, role, oauth=oauth
+                id, name, email, profile_image_url, role, oauth=oauth, db=db
             )
 
             db.commit()
@@ -112,16 +114,16 @@ class AuthsTable:
                 return None
 
     def authenticate_user(
-        self, email: str, verify_password: callable
+        self, email: str, verify_password: callable, db: Optional[Session] = None
     ) -> Optional[UserModel]:
         log.info(f"authenticate_user: {email}")
 
-        user = Users.get_user_by_email(email)
+        user = Users.get_user_by_email(email, db=db)
         if not user:
             return None
 
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 auth = db.query(Auth).filter_by(id=user.id, active=True).first()
                 if auth:
                     if verify_password(auth.password):
@@ -133,32 +135,32 @@ class AuthsTable:
         except Exception:
             return None
 
-    def authenticate_user_by_api_key(self, api_key: str) -> Optional[UserModel]:
+    def authenticate_user_by_api_key(self, api_key: str, db: Optional[Session] = None) -> Optional[UserModel]:
         log.info(f"authenticate_user_by_api_key: {api_key}")
         # if no api_key, return None
         if not api_key:
             return None
 
         try:
-            user = Users.get_user_by_api_key(api_key)
+            user = Users.get_user_by_api_key(api_key, db=db)
             return user if user else None
         except Exception:
             return False
 
-    def authenticate_user_by_email(self, email: str) -> Optional[UserModel]:
+    def authenticate_user_by_email(self, email: str, db: Optional[Session] = None) -> Optional[UserModel]:
         log.info(f"authenticate_user_by_email: {email}")
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 auth = db.query(Auth).filter_by(email=email, active=True).first()
                 if auth:
-                    user = Users.get_user_by_id(auth.id)
+                    user = Users.get_user_by_id(auth.id, db=db)
                     return user
         except Exception:
             return None
 
-    def update_user_password_by_id(self, id: str, new_password: str) -> bool:
+    def update_user_password_by_id(self, id: str, new_password: str, db: Optional[Session] = None) -> bool:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 result = (
                     db.query(Auth).filter_by(id=id).update({"password": new_password})
                 )
@@ -167,20 +169,20 @@ class AuthsTable:
         except Exception:
             return False
 
-    def update_email_by_id(self, id: str, email: str) -> bool:
+    def update_email_by_id(self, id: str, email: str, db: Optional[Session] = None) -> bool:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 result = db.query(Auth).filter_by(id=id).update({"email": email})
                 db.commit()
                 return True if result == 1 else False
         except Exception:
             return False
 
-    def delete_auth_by_id(self, id: str) -> bool:
+    def delete_auth_by_id(self, id: str, db: Optional[Session] = None) -> bool:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 # Delete User
-                result = Users.delete_user_by_id(id)
+                result = Users.delete_user_by_id(id, db=db)
 
                 if result:
                     db.query(Auth).filter_by(id=id).delete()