diff --git a/src/ohmyapi/builtin/auth/models.py b/src/ohmyapi/builtin/auth/models.py index 7628c4f..c0983c7 100644 --- a/src/ohmyapi/builtin/auth/models.py +++ b/src/ohmyapi/builtin/auth/models.py @@ -1,20 +1,22 @@ -from functools import wraps -from typing import Optional, List from ohmyapi.router import HTTPException from ohmyapi.db import Model, field, pre_save, pre_delete + +from functools import wraps +from typing import Optional, List from passlib.context import CryptContext from tortoise.contrib.pydantic import pydantic_queryset_creator +from uuid import UUID pwd_context = CryptContext(schemes=["argon2"], deprecated="auto") class Group(Model): - id: str = field.data.UUIDField(pk=True) + id: UUID = field.data.UUIDField(pk=True) name: str = field.CharField(max_length=42, index=True) class User(Model): - id: str = field.data.UUIDField(pk=True) + id: UUID = field.data.UUIDField(pk=True) email: str = field.CharField(max_length=255, unique=True, index=True) username: str = field.CharField(max_length=150, unique=True) password_hash: str = field.CharField(max_length=128) diff --git a/src/ohmyapi/db/model/model.py b/src/ohmyapi/db/model/model.py index eb0b504..e67b81e 100644 --- a/src/ohmyapi/db/model/model.py +++ b/src/ohmyapi/db/model/model.py @@ -1,6 +1,27 @@ +from pydantic_core import core_schema +from pydantic import GetCoreSchemaHandler from tortoise import fields as field from tortoise.models import Model as TortoiseModel from tortoise.contrib.pydantic import pydantic_model_creator, pydantic_queryset_creator +from uuid import UUID + + +def __uuid_schema_monkey_patch(cls, source_type, handler): + # Always treat UUID as string schema + return core_schema.no_info_after_validator_function( + # Accept UUID or str, always return UUID internally + lambda v: v if isinstance(v, UUID) else UUID(str(v)), + core_schema.union_schema([ + core_schema.str_schema(), + core_schema.is_instance_schema(UUID), + ]), + # But when serializing, always str() + serialization=core_schema.plain_serializer_function_ser_schema(str, when_used="always"), + ) + + +# Monkey-patch UUID +UUID.__get_pydantic_core_schema__ = classmethod(__uuid_schema_monkey_patch) class ModelMeta(type(TortoiseModel)):