refac: defer profile
This commit is contained in:
@@ -1,7 +1,7 @@
|
|||||||
import time
|
import time
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
from sqlalchemy.orm import Session
|
from sqlalchemy.orm import Session, defer
|
||||||
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
|
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
|
||||||
|
|
||||||
|
|
||||||
@@ -15,7 +15,7 @@ from open_webui.utils.misc import throttle
|
|||||||
from open_webui.utils.validate import validate_profile_image_url
|
from open_webui.utils.validate import validate_profile_image_url
|
||||||
|
|
||||||
|
|
||||||
from pydantic import BaseModel, ConfigDict, field_validator
|
from pydantic import BaseModel, ConfigDict, field_validator, model_validator
|
||||||
from sqlalchemy import (
|
from sqlalchemy import (
|
||||||
BigInteger,
|
BigInteger,
|
||||||
JSON,
|
JSON,
|
||||||
@@ -28,7 +28,7 @@ from sqlalchemy import (
|
|||||||
select,
|
select,
|
||||||
cast,
|
cast,
|
||||||
)
|
)
|
||||||
from sqlalchemy import or_, case
|
from sqlalchemy import or_, case, func
|
||||||
from sqlalchemy.dialects.postgresql import JSONB
|
from sqlalchemy.dialects.postgresql import JSONB
|
||||||
|
|
||||||
import datetime
|
import datetime
|
||||||
@@ -86,7 +86,7 @@ class UserModel(BaseModel):
|
|||||||
|
|
||||||
name: str
|
name: str
|
||||||
|
|
||||||
profile_image_url: str
|
profile_image_url: Optional[str] = None
|
||||||
profile_banner_image_url: Optional[str] = None
|
profile_banner_image_url: Optional[str] = None
|
||||||
|
|
||||||
bio: Optional[str] = None
|
bio: Optional[str] = None
|
||||||
@@ -110,6 +110,12 @@ class UserModel(BaseModel):
|
|||||||
|
|
||||||
model_config = ConfigDict(from_attributes=True)
|
model_config = ConfigDict(from_attributes=True)
|
||||||
|
|
||||||
|
@model_validator(mode="after")
|
||||||
|
def set_profile_image_url(self):
|
||||||
|
if not self.profile_image_url:
|
||||||
|
self.profile_image_url = f"/api/v1/users/{self.id}/profile/image"
|
||||||
|
return self
|
||||||
|
|
||||||
|
|
||||||
class UserStatusModel(UserModel):
|
class UserStatusModel(UserModel):
|
||||||
is_active: bool = False
|
is_active: bool = False
|
||||||
@@ -315,8 +321,12 @@ class UsersTable:
|
|||||||
) -> Optional[UserModel]:
|
) -> Optional[UserModel]:
|
||||||
try:
|
try:
|
||||||
with get_db_context(db) as db:
|
with get_db_context(db) as db:
|
||||||
user = db.query(User).filter_by(email=email).first()
|
user = (
|
||||||
return UserModel.model_validate(user)
|
db.query(User)
|
||||||
|
.filter(func.lower(User.email) == email.lower())
|
||||||
|
.first()
|
||||||
|
)
|
||||||
|
return UserModel.model_validate(user) if user else None
|
||||||
except Exception:
|
except Exception:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
@@ -350,7 +360,7 @@ class UsersTable:
|
|||||||
) -> dict:
|
) -> dict:
|
||||||
with get_db_context(db) as db:
|
with get_db_context(db) as db:
|
||||||
# Join GroupMember so we can order by group_id when requested
|
# Join GroupMember so we can order by group_id when requested
|
||||||
query = db.query(User)
|
query = db.query(User).options(defer(User.profile_image_url))
|
||||||
|
|
||||||
if filter:
|
if filter:
|
||||||
query_key = filter.get("query")
|
query_key = filter.get("query")
|
||||||
@@ -485,6 +495,7 @@ class UsersTable:
|
|||||||
with get_db_context(db) as db:
|
with get_db_context(db) as db:
|
||||||
users = (
|
users = (
|
||||||
db.query(User)
|
db.query(User)
|
||||||
|
.options(defer(User.profile_image_url))
|
||||||
.join(GroupMember, User.id == GroupMember.user_id)
|
.join(GroupMember, User.id == GroupMember.user_id)
|
||||||
.filter(GroupMember.group_id == group_id)
|
.filter(GroupMember.group_id == group_id)
|
||||||
.all()
|
.all()
|
||||||
@@ -495,7 +506,7 @@ class UsersTable:
|
|||||||
self, user_ids: list[str], db: Optional[Session] = None
|
self, user_ids: list[str], db: Optional[Session] = None
|
||||||
) -> list[UserStatusModel]:
|
) -> list[UserStatusModel]:
|
||||||
with get_db_context(db) as db:
|
with get_db_context(db) as db:
|
||||||
users = db.query(User).filter(User.id.in_(user_ids)).all()
|
users = db.query(User).options(defer(User.profile_image_url)).filter(User.id.in_(user_ids)).all()
|
||||||
return [UserModel.model_validate(user) for user in users]
|
return [UserModel.model_validate(user) for user in users]
|
||||||
|
|
||||||
def get_num_users(self, db: Optional[Session] = None) -> Optional[int]:
|
def get_num_users(self, db: Optional[Session] = None) -> Optional[int]:
|
||||||
|
|||||||
@@ -145,7 +145,7 @@ def create_session_response(
|
|||||||
"email": user.email,
|
"email": user.email,
|
||||||
"name": user.name,
|
"name": user.name,
|
||||||
"role": user.role,
|
"role": user.role,
|
||||||
"profile_image_url": user.profile_image_url,
|
"profile_image_url": f"/api/v1/users/{user.id}/profile/image",
|
||||||
"permissions": user_permissions,
|
"permissions": user_permissions,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -926,7 +926,7 @@ async def add_user(
|
|||||||
"email": user.email,
|
"email": user.email,
|
||||||
"name": user.name,
|
"name": user.name,
|
||||||
"role": user.role,
|
"role": user.role,
|
||||||
"profile_image_url": user.profile_image_url,
|
"profile_image_url": f"/api/v1/users/{user.id}/profile/image",
|
||||||
}
|
}
|
||||||
else:
|
else:
|
||||||
raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_USER_ERROR)
|
raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_USER_ERROR)
|
||||||
|
|||||||
@@ -314,7 +314,13 @@ async def connect(sid, environ, auth):
|
|||||||
|
|
||||||
if user:
|
if user:
|
||||||
SESSION_POOL[sid] = user.model_dump(
|
SESSION_POOL[sid] = user.model_dump(
|
||||||
exclude=["date_of_birth", "bio", "gender"]
|
exclude=[
|
||||||
|
"profile_image_url",
|
||||||
|
"profile_banner_image_url",
|
||||||
|
"date_of_birth",
|
||||||
|
"bio",
|
||||||
|
"gender",
|
||||||
|
]
|
||||||
)
|
)
|
||||||
await sio.enter_room(sid, f"user:{user.id}")
|
await sio.enter_room(sid, f"user:{user.id}")
|
||||||
|
|
||||||
|
|||||||
@@ -12,7 +12,7 @@ def _assert_user(data, id, **kwargs):
|
|||||||
comparison_data = {
|
comparison_data = {
|
||||||
"name": f"user {id}",
|
"name": f"user {id}",
|
||||||
"email": f"user{id}@openwebui.com",
|
"email": f"user{id}@openwebui.com",
|
||||||
"profile_image_url": f"/user{id}.png",
|
"profile_image_url": f"/api/v1/users/{id}/profile/image",
|
||||||
"role": "user",
|
"role": "user",
|
||||||
**kwargs,
|
**kwargs,
|
||||||
}
|
}
|
||||||
@@ -150,7 +150,7 @@ class TestUsers(AbstractPostgresTest):
|
|||||||
role="admin",
|
role="admin",
|
||||||
name="user 2 updated",
|
name="user 2 updated",
|
||||||
email="user2-updated@openwebui.com",
|
email="user2-updated@openwebui.com",
|
||||||
profile_image_url="/user2-updated.png",
|
profile_image_url=f"/api/v1/users/2/profile/image",
|
||||||
)
|
)
|
||||||
|
|
||||||
# Delete user by id
|
# Delete user by id
|
||||||
|
|||||||
Reference in New Issue
Block a user