feat(sqlalchemy): Replace peewee with sqlalchemy

backend/apps/webui/models/users.py

@@ -1,11 +1,13 @@
-from pydantic import BaseModel, ConfigDict
-from peewee import *
-from playhouse.shortcuts import model_to_dict
+from pydantic import BaseModel, ConfigDict, parse_obj_as
 from typing import List, Union, Optional
 import time
+
+from sqlalchemy import String, Column, BigInteger, Text
+from sqlalchemy.orm import Session
+
 from utils.misc import get_gravatar_url
 
-from apps.webui.internal.db import DB, JSONField
+from apps.webui.internal.db import Base, JSONField
 from apps.webui.models.chats import Chats
 
 ####################
@@ -13,25 +15,24 @@ from apps.webui.models.chats import Chats
 ####################
 
 
-class User(Model):
-    id = CharField(unique=True)
-    name = CharField()
-    email = CharField()
-    role = CharField()
-    profile_image_url = TextField()
+class User(Base):
+    __tablename__ = "user"
 
-    last_active_at = BigIntegerField()
-    updated_at = BigIntegerField()
-    created_at = BigIntegerField()
+    id = Column(String, primary_key=True)
+    name = Column(String)
+    email = Column(String)
+    role = Column(String)
+    profile_image_url = Column(String)
 
-    api_key = CharField(null=True, unique=True)
-    settings = JSONField(null=True)
-    info = JSONField(null=True)
+    last_active_at = Column(BigInteger)
+    updated_at = Column(BigInteger)
+    created_at = Column(BigInteger)
 
-    oauth_sub = TextField(null=True, unique=True)
+    api_key = Column(String, nullable=True, unique=True)
+    settings = Column(JSONField, nullable=True)
+    info = Column(JSONField, nullable=True)
 
-    class Meta:
-        database = DB
+    oauth_sub = Column(Text, unique=True)
 
 
 class UserSettings(BaseModel):
@@ -41,6 +42,8 @@ class UserSettings(BaseModel):
 
 
 class UserModel(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
     id: str
     name: str
     email: str
@@ -76,12 +79,10 @@ class UserUpdateForm(BaseModel):
 
 
 class UsersTable:
-    def __init__(self, db):
-        self.db = db
-        self.db.create_tables([User])
 
     def insert_new_user(
         self,
+        db: Session,
         id: str,
         name: str,
         email: str,
@@ -102,30 +103,33 @@ class UsersTable:
                 "oauth_sub": oauth_sub,
             }
         )
-        result = User.create(**user.model_dump())
+        result = User(**user.model_dump())
+        db.add(result)
+        db.commit()
+        db.refresh(result)
         if result:
             return user
         else:
             return None
 
-    def get_user_by_id(self, id: str) -> Optional[UserModel]:
+    def get_user_by_id(self, db: Session, id: str) -> Optional[UserModel]:
         try:
-            user = User.get(User.id == id)
-            return UserModel(**model_to_dict(user))
+            user = db.query(User).filter_by(id=id).first()
+            return UserModel.model_validate(user)
+        except Exception as e:
+            return None
+
+    def get_user_by_api_key(self, db: Session, api_key: str) -> Optional[UserModel]:
+        try:
+            user = db.query(User).filter_by(api_key=api_key).first()
+            return UserModel.model_validate(user)
         except:
             return None
 
-    def get_user_by_api_key(self, api_key: str) -> Optional[UserModel]:
+    def get_user_by_email(self, db: Session, email: str) -> Optional[UserModel]:
         try:
-            user = User.get(User.api_key == api_key)
-            return UserModel(**model_to_dict(user))
-        except:
-            return None
-
-    def get_user_by_email(self, email: str) -> Optional[UserModel]:
-        try:
-            user = User.get(User.email == email)
-            return UserModel(**model_to_dict(user))
+            user = db.query(User).filter_by(email=email).first()
+            return UserModel.model_validate(user)
         except:
             return None
 
@@ -136,88 +140,94 @@ class UsersTable:
         except:
             return None
 
-    def get_users(self, skip: int = 0, limit: int = 50) -> List[UserModel]:
-        return [
-            UserModel(**model_to_dict(user))
-            for user in User.select()
-            # .limit(limit).offset(skip)
-        ]
+    def get_users(self, db: Session, skip: int = 0, limit: int = 50) -> List[UserModel]:
+        users = (
+            db.query(User)
+            # .offset(skip).limit(limit)
+            .all()
+        )
+        return [UserModel.model_validate(user) for user in users]
 
-    def get_num_users(self) -> Optional[int]:
-        return User.select().count()
+    def get_num_users(self, db: Session) -> Optional[int]:
+        return db.query(User).count()
 
-    def get_first_user(self) -> UserModel:
+    def get_first_user(self, db: Session) -> UserModel:
         try:
-            user = User.select().order_by(User.created_at).first()
-            return UserModel(**model_to_dict(user))
+            user = db.query(User).order_by(User.created_at).first()
+            return UserModel.model_validate(user)
         except:
             return None
 
-    def update_user_role_by_id(self, id: str, role: str) -> Optional[UserModel]:
+    def update_user_role_by_id(
+        self, db: Session, id: str, role: str
+    ) -> Optional[UserModel]:
         try:
-            query = User.update(role=role).where(User.id == id)
-            query.execute()
+            db.query(User).filter_by(id=id).update({"role": role})
+            db.commit()
 
-            user = User.get(User.id == id)
-            return UserModel(**model_to_dict(user))
+            user = db.query(User).filter_by(id=id).first()
+            return UserModel.model_validate(user)
         except:
             return None
 
     def update_user_profile_image_url_by_id(
-        self, id: str, profile_image_url: str
+        self, db: Session, id: str, profile_image_url: str
     ) -> Optional[UserModel]:
         try:
-            query = User.update(profile_image_url=profile_image_url).where(
-                User.id == id
+            db.query(User).filter_by(id=id).update(
+                {"profile_image_url": profile_image_url}
             )
-            query.execute()
+            db.commit()
 
-            user = User.get(User.id == id)
-            return UserModel(**model_to_dict(user))
+            user = db.query(User).filter_by(id=id).first()
+            return UserModel.model_validate(user)
         except:
             return None
 
-    def update_user_last_active_by_id(self, id: str) -> Optional[UserModel]:
+    def update_user_last_active_by_id(
+        self, db: Session, id: str
+    ) -> Optional[UserModel]:
         try:
-            query = User.update(last_active_at=int(time.time())).where(User.id == id)
-            query.execute()
+            db.query(User).filter_by(id=id).update({"last_active_at": int(time.time())})
 
-            user = User.get(User.id == id)
-            return UserModel(**model_to_dict(user))
+            user = db.query(User).filter_by(id=id).first()
+            return UserModel.model_validate(user)
         except:
             return None
 
     def update_user_oauth_sub_by_id(
-        self, id: str, oauth_sub: str
+        self, db: Session, id: str, oauth_sub: str
     ) -> Optional[UserModel]:
         try:
-            query = User.update(oauth_sub=oauth_sub).where(User.id == id)
-            query.execute()
+            db.query(User).filter_by(id=id).update({"oauth_sub": oauth_sub})
 
-            user = User.get(User.id == id)
-            return UserModel(**model_to_dict(user))
+            user = db.query(User).filter_by(id=id).first()
+            return UserModel.model_validate(user)
         except:
             return None
 
-    def update_user_by_id(self, id: str, updated: dict) -> Optional[UserModel]:
+    def update_user_by_id(
+        self, db: Session, id: str, updated: dict
+    ) -> Optional[UserModel]:
         try:
-            query = User.update(**updated).where(User.id == id)
-            query.execute()
+            db.query(User).filter_by(id=id).update(updated)
+            db.commit()
 
-            user = User.get(User.id == id)
-            return UserModel(**model_to_dict(user))
-        except:
+            user = db.query(User).filter_by(id=id).first()
+            return UserModel.model_validate(user)
+            # return UserModel(**user.dict())
+        except Exception as e:
             return None
 
-    def delete_user_by_id(self, id: str) -> bool:
+    def delete_user_by_id(self, db: Session, id: str) -> bool:
         try:
             # Delete User Chats
-            result = Chats.delete_chats_by_user_id(id)
+            result = Chats.delete_chats_by_user_id(db, id)
 
             if result:
                 # Delete User
-                query = User.delete().where(User.id == id)
-                query.execute()  # Remove the rows, return number of rows removed.
+                db.query(User).filter_by(id=id).delete()
+                db.commit()
 
                 return True
             else:
@@ -225,21 +235,20 @@ class UsersTable:
         except:
             return False
 
-    def update_user_api_key_by_id(self, id: str, api_key: str) -> str:
+    def update_user_api_key_by_id(self, db: Session, id: str, api_key: str) -> str:
         try:
-            query = User.update(api_key=api_key).where(User.id == id)
-            result = query.execute()
-
+            result = db.query(User).filter_by(id=id).update({"api_key": api_key})
+            db.commit()
             return True if result == 1 else False
         except:
             return False
 
-    def get_user_api_key_by_id(self, id: str) -> Optional[str]:
+    def get_user_api_key_by_id(self, db: Session, id: str) -> Optional[str]:
         try:
-            user = User.get(User.id == id)
+            user = db.query(User).filter_by(id=id).first()
             return user.api_key
-        except:
+        except Exception as e:
             return None
 
 
-Users = UsersTable(DB)
+Users = UsersTable()