Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy J. Baek
2024-07-03 23:32:39 -07:00
parent 1b65df3acc
commit 864646094e
11 changed files with 789 additions and 616 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/webui/models/users.py
+91 -72
View File
@@ -6,7 +6,7 @@ from sqlalchemy import String, Column, BigInteger, Text
from utils.misc import get_gravatar_url
from apps.webui.internal.db import Base, JSONField, Session
from apps.webui.internal.db import Base, JSONField, Session, get_db
from apps.webui.models.chats import Chats
####################
@@ -88,81 +88,92 @@ class UsersTable:
role: str = "pending",
oauth_sub: Optional[str] = None,
) -> Optional[UserModel]:
user = UserModel(
**{
"id": id,
"name": name,
"email": email,
"role": role,
"profile_image_url": profile_image_url,
"last_active_at": int(time.time()),
"created_at": int(time.time()),
"updated_at": int(time.time()),
"oauth_sub": oauth_sub,
}
)
result = User(**user.model_dump())
Session.add(result)
Session.commit()
Session.refresh(result)
if result:
return user
else:
return None
with get_db() as db:
user = UserModel(
**{
"id": id,
"name": name,
"email": email,
"role": role,
"profile_image_url": profile_image_url,
"last_active_at": int(time.time()),
"created_at": int(time.time()),
"updated_at": int(time.time()),
"oauth_sub": oauth_sub,
}
)
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]:
try:
user = Session.query(User).filter_by(id=id).first()
return UserModel.model_validate(user)
with get_db() as db:
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, api_key: str) -> Optional[UserModel]:
try:
user = Session.query(User).filter_by(api_key=api_key).first()
return UserModel.model_validate(user)
with get_db() as db:
user = db.query(User).filter_by(api_key=api_key).first()
return UserModel.model_validate(user)
except:
return None
def get_user_by_email(self, email: str) -> Optional[UserModel]:
try:
user = Session.query(User).filter_by(email=email).first()
return UserModel.model_validate(user)
with get_db() as db:
user = db.query(User).filter_by(email=email).first()
return UserModel.model_validate(user)
except:
return None
def get_user_by_oauth_sub(self, sub: str) -> Optional[UserModel]:
try:
user = Session.query(User).filter_by(oauth_sub=sub).first()
return UserModel.model_validate(user)
with get_db() as db:
user = db.query(User).filter_by(oauth_sub=sub).first()
return UserModel.model_validate(user)
except:
return None
def get_users(self, skip: int = 0, limit: int = 50) -> List[UserModel]:
users = (
Session.query(User)
# .offset(skip).limit(limit)
.all()
)
return [UserModel.model_validate(user) for user in users]
with get_db() as db:
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 Session.query(User).count()
with get_db() as db:
return db.query(User).count()
def get_first_user(self) -> UserModel:
try:
user = Session.query(User).order_by(User.created_at).first()
return UserModel.model_validate(user)
with get_db() as db:
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]:
try:
Session.query(User).filter_by(id=id).update({"role": role})
Session.commit()
user = Session.query(User).filter_by(id=id).first()
return UserModel.model_validate(user)
with get_db() as db:
db.query(User).filter_by(id=id).update({"role": role})
db.commit()
user = db.query(User).filter_by(id=id).first()
return UserModel.model_validate(user)
except:
return None
@@ -170,25 +181,28 @@ class UsersTable:
self, id: str, profile_image_url: str
) -> Optional[UserModel]:
try:
Session.query(User).filter_by(id=id).update(
{"profile_image_url": profile_image_url}
)
Session.commit()
with get_db() as db:
db.query(User).filter_by(id=id).update(
{"profile_image_url": profile_image_url}
)
db.commit()
user = Session.query(User).filter_by(id=id).first()
return UserModel.model_validate(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]:
try:
Session.query(User).filter_by(id=id).update(
{"last_active_at": int(time.time())}
)
Session.commit()
with get_db() as db:
user = Session.query(User).filter_by(id=id).first()
return UserModel.model_validate(user)
db.query(User).filter_by(id=id).update(
{"last_active_at": int(time.time())}
)
db.commit()
user = db.query(User).filter_by(id=id).first()
return UserModel.model_validate(user)
except:
return None
@@ -196,21 +210,23 @@ class UsersTable:
self, id: str, oauth_sub: str
) -> Optional[UserModel]:
try:
Session.query(User).filter_by(id=id).update({"oauth_sub": oauth_sub})
with get_db() as db:
db.query(User).filter_by(id=id).update({"oauth_sub": oauth_sub})
user = Session.query(User).filter_by(id=id).first()
return UserModel.model_validate(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]:
try:
Session.query(User).filter_by(id=id).update(updated)
Session.commit()
with get_db() as db:
db.query(User).filter_by(id=id).update(updated)
db.commit()
user = Session.query(User).filter_by(id=id).first()
return UserModel.model_validate(user)
# return UserModel(**user.dict())
user = db.query(User).filter_by(id=id).first()
return UserModel.model_validate(user)
# return UserModel(**user.dict())
except Exception as e:
return None
@@ -220,9 +236,10 @@ class UsersTable:
result = Chats.delete_chats_by_user_id(id)
if result:
# Delete User
Session.query(User).filter_by(id=id).delete()
Session.commit()
with get_db() as db:
# Delete User
db.query(User).filter_by(id=id).delete()
db.commit()
return True
else:
@@ -232,16 +249,18 @@ class UsersTable:
def update_user_api_key_by_id(self, id: str, api_key: str) -> str:
try:
result = Session.query(User).filter_by(id=id).update({"api_key": api_key})
Session.commit()
return True if result == 1 else False
with get_db() as db:
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]:
try:
user = Session.query(User).filter_by(id=id).first()
return user.api_key
with get_db() as db:
user = db.query(User).filter_by(id=id).first()
return user.api_key
except Exception as e:
return None