105 lines
No EOL
3.3 KiB
Python
105 lines
No EOL
3.3 KiB
Python
from pydantic import BaseModel, Field, EmailStr
|
|
from typing import Optional, List, Dict
|
|
from datetime import datetime
|
|
from bson import ObjectId
|
|
from enum import Enum
|
|
|
|
class UserRole(str, Enum):
|
|
ADMIN = "admin"
|
|
USER = "user"
|
|
|
|
class AuthMethod(str, Enum):
|
|
LOCAL = "local"
|
|
SSO = "sso"
|
|
|
|
class PyObjectId(ObjectId):
|
|
@classmethod
|
|
def __get_pydantic_core_schema__(cls, _source_type, _handler):
|
|
from pydantic_core import core_schema
|
|
|
|
def validate_from_str(input_value: str) -> ObjectId:
|
|
return ObjectId(input_value)
|
|
|
|
def validate_from_objectid(input_value: ObjectId) -> ObjectId:
|
|
return input_value
|
|
|
|
return core_schema.union_schema([
|
|
core_schema.is_instance_schema(ObjectId),
|
|
core_schema.no_info_plain_validator_function(
|
|
validate_from_str,
|
|
serialization=core_schema.to_string_ser_schema(),
|
|
),
|
|
])
|
|
|
|
@classmethod
|
|
def __get_pydantic_json_schema__(cls, core_schema, handler):
|
|
json_schema = handler(core_schema)
|
|
json_schema.update(type="string")
|
|
return json_schema
|
|
|
|
class UserBase(BaseModel):
|
|
email: EmailStr
|
|
role: UserRole = UserRole.USER
|
|
is_active: bool = True
|
|
auth_method: AuthMethod = AuthMethod.LOCAL
|
|
sso_provider: Optional[str] = None
|
|
sso_user_id: Optional[str] = None
|
|
sso_email: Optional[str] = None
|
|
sso_name: Optional[str] = None
|
|
sso_attributes: Optional[Dict] = None
|
|
last_sso_login: Optional[datetime] = None
|
|
created_at: Optional[datetime] = None
|
|
updated_at: Optional[datetime] = None
|
|
|
|
class UserCreate(UserBase):
|
|
password: str
|
|
|
|
class UserUpdate(BaseModel):
|
|
email: Optional[EmailStr] = None
|
|
role: Optional[UserRole] = None
|
|
is_active: Optional[bool] = None
|
|
password: Optional[str] = None
|
|
auth_method: Optional[AuthMethod] = None
|
|
sso_provider: Optional[str] = None
|
|
sso_user_id: Optional[str] = None
|
|
sso_email: Optional[str] = None
|
|
sso_name: Optional[str] = None
|
|
sso_attributes: Optional[Dict] = None
|
|
last_sso_login: Optional[datetime] = None
|
|
|
|
class UserInDB(UserBase):
|
|
id: PyObjectId = Field(default_factory=PyObjectId, alias="_id")
|
|
hashed_password: Optional[str] = None # Optional for SSO users
|
|
index_access: List[str] = Field(default_factory=list) # List of index IDs user has access to
|
|
|
|
class Config:
|
|
populate_by_name = True
|
|
arbitrary_types_allowed = True
|
|
json_encoders = {ObjectId: str}
|
|
|
|
class User(UserBase):
|
|
id: PyObjectId = Field(default_factory=PyObjectId, alias="_id")
|
|
index_access: List[str] = Field(default_factory=list)
|
|
|
|
class Config:
|
|
populate_by_name = True
|
|
arbitrary_types_allowed = True
|
|
json_encoders = {ObjectId: str}
|
|
|
|
class UserResponse(BaseModel):
|
|
id: PyObjectId = Field(alias="_id")
|
|
email: EmailStr
|
|
role: UserRole
|
|
is_active: bool
|
|
auth_method: AuthMethod
|
|
sso_provider: Optional[str] = None
|
|
sso_name: Optional[str] = None
|
|
last_sso_login: Optional[datetime] = None
|
|
index_access: List[str]
|
|
created_at: Optional[datetime] = None
|
|
updated_at: Optional[datetime] = None
|
|
|
|
class Config:
|
|
populate_by_name = True
|
|
arbitrary_types_allowed = True
|
|
json_encoders = {ObjectId: str} |