oidc-fastapi-test/src/oidc_test/models.py

from functools import cached_property
from typing import Self

from pydantic import BaseModel, EmailStr, AnyHttpUrl, Field, computed_field
from authlib.integrations.starlette_client.apps import StarletteOAuth2App

# from app.models import User


class Role(BaseModel, extra="ignore"):
    name: str


class UserBase(BaseModel, extra="ignore"):

    id: str | None = None
    name: str
    email: EmailStr | None = None
    picture: AnyHttpUrl | None = None
    roles: list[Role] = []


class User(UserBase):
    class Config:
        arbitrary_types_allowed = True

    sub: str = Field(
        description="""subject id of the user given by the oidc provider,
                            also the key for the database 'table'""",
    )
    userinfo: dict = {}
    oidc_provider: StarletteOAuth2App | None = None

    @classmethod
    def from_auth(cls, userinfo: dict, oidc_provider: StarletteOAuth2App) -> Self:
        user = cls(**userinfo)
        user.userinfo = userinfo
        user.oidc_provider = oidc_provider
        # Add roles if they are provided in the token
        if raw_ra := userinfo.get("realm_access"):
            if raw_roles := raw_ra.get("roles"):
                user.roles = [Role(name=raw_role) for raw_role in raw_roles]
        return user

    @computed_field
    @cached_property
    def roles_as_set(self) -> set[str]:
        return set([role.name for role in self.roles])