feat(sqlalchemy): Replace peewee with sqlalchemy

backend/apps/webui/models/auths.py

@@ -1,14 +1,14 @@
 from pydantic import BaseModel
-from typing import List, Union, Optional
-import time
+from typing import Optional
 import uuid
 import logging
-from peewee import *
+from sqlalchemy import String, Column, Boolean
+from sqlalchemy.orm import Session
 
 from apps.webui.models.users import UserModel, Users
 from utils.utils import verify_password
 
-from apps.webui.internal.db import DB
+from apps.webui.internal.db import Base
 
 from config import SRC_LOG_LEVELS
 
@@ -20,14 +20,13 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
 ####################
 
 
-class Auth(Model):
-    id = CharField(unique=True)
-    email = CharField()
-    password = TextField()
-    active = BooleanField()
+class Auth(Base):
+    __tablename__ = "auth"
 
-    class Meta:
-        database = DB
+    id = Column(String, primary_key=True)
+    email = Column(String)
+    password = Column(String)
+    active = Column(Boolean)
 
 
 class AuthModel(BaseModel):
@@ -94,12 +93,10 @@ class AddUserForm(SignupForm):
 
 
 class AuthsTable:
-    def __init__(self, db):
-        self.db = db
-        self.db.create_tables([Auth])
 
     def insert_new_auth(
         self,
+        db: Session,
         email: str,
         password: str,
         name: str,
@@ -114,24 +111,30 @@ class AuthsTable:
         auth = AuthModel(
             **{"id": id, "email": email, "password": password, "active": True}
         )
-        result = Auth.create(**auth.model_dump())
+        result = Auth(**auth.model_dump())
+        db.add(result)
 
         user = Users.insert_new_user(
-            id, name, email, profile_image_url, role, oauth_sub
+            db, id, name, email, profile_image_url, role, oauth_sub
         )
 
+        db.commit()
+        db.refresh(result)
+
         if result and user:
             return user
         else:
             return None
 
-    def authenticate_user(self, email: str, password: str) -> Optional[UserModel]:
+    def authenticate_user(
+        self, db: Session, email: str, password: str
+    ) -> Optional[UserModel]:
         log.info(f"authenticate_user: {email}")
         try:
-            auth = Auth.get(Auth.email == email, Auth.active == True)
+            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)
+                    user = Users.get_user_by_id(db, auth.id)
                     return user
                 else:
                     return None
@@ -140,55 +143,55 @@ class AuthsTable:
         except:
             return None
 
-    def authenticate_user_by_api_key(self, api_key: str) -> Optional[UserModel]:
+    def authenticate_user_by_api_key(
+        self, db: Session, 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
 
         try:
-            user = Users.get_user_by_api_key(api_key)
+            user = Users.get_user_by_api_key(db, api_key)
             return user if user else None
         except:
             return False
 
-    def authenticate_user_by_trusted_header(self, email: str) -> Optional[UserModel]:
+    def authenticate_user_by_trusted_header(
+        self, db: Session, email: str
+    ) -> Optional[UserModel]:
         log.info(f"authenticate_user_by_trusted_header: {email}")
         try:
-            auth = Auth.get(Auth.email == email, Auth.active == True)
+            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, id: str, new_password: str) -> bool:
+    def update_user_password_by_id(
+        self, db: Session, id: str, new_password: str
+    ) -> bool:
         try:
-            query = Auth.update(password=new_password).where(Auth.id == id)
-            result = query.execute()
-
+            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:
+    def update_email_by_id(self, db: Session, id: str, email: str) -> bool:
         try:
-            query = Auth.update(email=email).where(Auth.id == id)
-            result = query.execute()
-
+            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:
+    def delete_auth_by_id(self, db: Session, id: str) -> bool:
         try:
             # Delete User
-            result = Users.delete_user_by_id(id)
+            result = Users.delete_user_by_id(db, id)
 
             if result:
-                # Delete Auth
-                query = Auth.delete().where(Auth.id == id)
-                query.execute()  # Remove the rows, return number of rows removed.
+                db.query(Auth).filter_by(id=id).delete()
 
                 return True
             else:
@@ -197,4 +200,4 @@ class AuthsTable:
             return False
 
 
-Auths = AuthsTable(DB)
+Auths = AuthsTable()