feat(sqlalchemy): remove session reference from router

open-webui · tjbck · Jul 6, 2024 · Jun 18, 2024 · Jun 21, 2024 · Jun 24, 2024
commit bee835cb65a8b3feba6824d2e6c9378b95f6e990
diff --git a/backend/apps/ollama/main.py b/backend/apps/ollama/main.py
@@ -31,7 31,6 @@
 
 from starlette.background import BackgroundTask
 
-from apps.webui.internal.db import get_db
 from apps.webui.models.models import Models
 from apps.webui.models.users import Users
 from constants import ERROR_MESSAGES
@@ -712,7 711,6 @@ async def generate_chat_completion(
  form_data: GenerateChatCompletionForm,
  url_idx: Optional[int] = None,
  user=Depends(get_verified_user),
- db=Depends(get_db),
 ):
 
  log.debug(
@@ -726,7 724,7 @@ async def generate_chat_completion(
  }
 
  model_id = form_data.model
- model_info = Models.get_model_by_id(db, model_id)
  model_info = Models.get_model_by_id(model_id)
 
  if model_info:
  if model_info.base_model_id:
@@ -885,7 883,6 @@ async def generate_openai_chat_completion(
  form_data: dict,
  url_idx: Optional[int] = None,
  user=Depends(get_verified_user),
- db=Depends(get_db),
 ):
  form_data = OpenAIChatCompletionForm(**form_data)
 
@@ -894,7 891,7 @@ async def generate_openai_chat_completion(
  }
 
  model_id = form_data.model
- model_info = Models.get_model_by_id(db, model_id)
  model_info = Models.get_model_by_id(model_id)
 
  if model_info:
  if model_info.base_model_id:

diff --git a/backend/apps/openai/main.py b/backend/apps/openai/main.py
@@ -11,7 11,6 @@
 from pydantic import BaseModel
 from starlette.background import BackgroundTask
 
-from apps.webui.internal.db import get_db
 from apps.webui.models.models import Models
 from apps.webui.models.users import Users
 from constants import ERROR_MESSAGES
@@ -354,13 353,12 @@ async def generate_chat_completion(
  form_data: dict,
  url_idx: Optional[int] = None,
  user=Depends(get_verified_user),
- db=Depends(get_db),
 ):
  idx = 0
  payload = {**form_data}
 
  model_id = form_data.get("model")
- model_info = Models.get_model_by_id(db, model_id)
  model_info = Models.get_model_by_id(model_id)
 
  if model_info:
  if model_info.base_model_id:

diff --git a/backend/apps/webui/internal/db.py b/backend/apps/webui/internal/db.py
@@ -1,6 1,7 @@
 import os
 import logging
 import json
 from contextlib import contextmanager
 from typing import Optional, Any
 from typing_extensions import Self
 
@@ -52,17 53,17 @@ def python_value(self, value):
  )
 else:
  engine = create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True)
-SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine, expire_on_commit=False)
 Base = declarative_base()
 
 
-def get_db():
 @contextmanager
 def get_session():
  db = SessionLocal()
  try:
  yield db
  db.commit()
  except Exception as e:
  db.rollback()
  raise e
- finally:
- db.close()
 
diff --git a/backend/apps/webui/main.py b/backend/apps/webui/main.py
@@ -114,8 114,8 @@ async def get_status():
  }
 
 
-async def get_pipe_models(db: Session):
- pipes = Functions.get_functions_by_type(db, "pipe", active_only=True)
 async def get_pipe_models():
  pipes = Functions.get_functions_by_type("pipe", active_only=True)
  pipe_models = []
 
  for pipe in pipes:

diff --git a/backend/apps/webui/models/auths.py b/backend/apps/webui/models/auths.py
@@ -8,7 8,7 @@
 from apps.webui.models.users import UserModel, Users
 from utils.utils import verify_password
 
-from apps.webui.internal.db import Base
 from apps.webui.internal.db import Base, get_session
 
 from config import SRC_LOG_LEVELS
 
@@ -96,108 96,114 @@ class AuthsTable:
 
  def insert_new_auth(
  self,
- db: Session,
  email: str,
  password: str,
  name: str,
  profile_image_url: str = "/user.png",
  role: str = "pending",
  oauth_sub: Optional[str] = None,
  ) -> Optional[UserModel]:
- log.info("insert_new_auth")
  with get_session() as db:
  log.info("insert_new_auth")
 
- id = str(uuid.uuid4())
   id = str(uuid.uuid4())
 
- auth = AuthModel(
- **{"id": id, "email": email, "password": password, "active": True}
- )
- result = Auth(**auth.model_dump())
- db.add(result)
   auth = AuthModel(
   **{"id": id, "email": email, "password": password, "active": True}
   )
   result = Auth(**auth.model_dump())
   db.add(result)
 
- user = Users.insert_new_user(
- db, id, name, email, profile_image_url, role, oauth_sub
- )
   user = Users.insert_new_user(
    id, name, email, profile_image_url, role, oauth_sub
   )
 
- db.commit()
- db.refresh(result)
   db.commit()
   db.refresh(result)
 
- if result and user:
- return user
- else:
- return None
   if result and user:
   return user
   else:
   return None
 
  def authenticate_user(
- self, db: Session, email: str, password: str
  self, email: str, password: str
  ) -> Optional[UserModel]:
  log.info(f"authenticate_user: {email}")
- try:
- auth = db.query(Auth).filter_by(email=email, active=True).first()
- if auth:
- if verify_password(password, auth.password):
- user = Users.get_user_by_id(db, auth.id)
- return user
  with get_session() as db:
  try:
  auth = db.query(Auth).filter_by(email=email, active=True).first()
  if auth:
  if verify_password(password, auth.password):
  user = Users.get_user_by_id(auth.id)
  return user
  else:
  return None
  else:
  return None
- else:
  except:
  return None
- except:
- return None
 
  def authenticate_user_by_api_key(
- self, db: Session, api_key: str
  self, api_key: str
  ) -> Optional[UserModel]:
  log.info(f"authenticate_user_by_api_key: {api_key}")
- # if no api_key, return None
- if not api_key:
- return None
  with get_session() as db:
  # if no api_key, return None
  if not api_key:
  return None
 
- try:
- user = Users.get_user_by_api_key(db, api_key)
- return user if user else None
- except:
- return False
   try:
   user = Users.get_user_by_api_key(api_key)
   return user if user else None
   except:
   return False
 
  def authenticate_user_by_trusted_header(
- self, db: Session, email: str
  self, email: str
  ) -> Optional[UserModel]:
  log.info(f"authenticate_user_by_trusted_header: {email}")
- try:
- auth = db.query(Auth).filter(email=email, active=True).first()
- if auth:
- user = Users.get_user_by_id(auth.id)
- return user
- except:
- return None
  with get_session() as db:
  try:
  auth = db.query(Auth).filter(email=email, active=True).first()
  if auth:
  user = Users.get_user_by_id(auth.id)
  return user
  except:
  return None
 
  def update_user_password_by_id(
- self, db: Session, id: str, new_password: str
  self, id: str, new_password: str
  ) -> bool:
- try:
- result = db.query(Auth).filter_by(id=id).update({"password": new_password})
- return True if result == 1 else False
- except:
- return False
-
- def update_email_by_id(self, db: Session, id: str, email: str) -> bool:
- try:
- result = db.query(Auth).filter_by(id=id).update({"email": email})
- return True if result == 1 else False
- except:
- return False
-
- def delete_auth_by_id(self, db: Session, id: str) -> bool:
- try:
- # Delete User
- result = Users.delete_user_by_id(db, id)
-
- if result:
- db.query(Auth).filter_by(id=id).delete()
-
- return True
- else:
  with get_session() as db:
  try:
  result = db.query(Auth).filter_by(id=id).update({"password": new_password})
  return True if result == 1 else False
  except:
  return False
 
  def update_email_by_id(self, id: str, email: str) -> bool:
  with get_session() as db:
  try:
  result = db.query(Auth).filter_by(id=id).update({"email": email})
  return True if result == 1 else False
  except:
  return False
 
  def delete_auth_by_id(self, id: str) -> bool:
  with get_session() as db:
  try:
  # Delete User
  result = Users.delete_user_by_id(id)
 
  if result:
  db.query(Auth).filter_by(id=id).delete()
 
  return True
  else:
  return False
  except:
  return False
- except:
- return False
 
 
 Auths = AuthsTable()