Add support for safe access to PK int autotypes

This commit is contained in:
Esteban Maya Cadavid 2024-07-01 17:28:11 -05:00
parent 3b889e09f7
commit 76d72cd32e
2 changed files with 65 additions and 0 deletions

View File

@ -219,6 +219,33 @@ if IS_PYDANTIC_V2:
) -> Optional[AbstractSet[str]]: # pragma: no cover
return None
def validate_access_primary_key_autotype(
self: InstanceOrType["SQLModel"], name: str, value: Any
) -> None:
"""
Pydantic v2
Validates if the attribute being accessed is a primary key with an auto type and has not been set.
Args:
self (InstanceOrType["SQLModel"]): The instance or type of SQLModel.
name (str): The name of the attribute being accessed.
value (Any): The value of the attribute being accessed.
Raises:
ValueError: If the attribute is a primary key with an auto type and has not been set.
Returns:
None
"""
if name != "model_fields":
model_fields = object.__getattribute__(self, "model_fields")
field = model_fields.get(name)
if field is not None and isinstance(field, FieldInfo):
if field.primary_key and field.annotation is int and value is None:
raise ValueError(
f"Primary key attribute '{name}' has not been set, please commit() it first."
)
def sqlmodel_table_construct(
*,
self_instance: _TSQLModel,
@ -499,6 +526,37 @@ else:
return keys
def validate_access_primary_key_autotype(
self: InstanceOrType["SQLModel"], name: str, value: Any
) -> None:
"""
Pydantic v1
Validates if the attribute being accessed is a primary key with an auto type and has not been set.
Args:
self (InstanceOrType["SQLModel"]): The instance or type of SQLModel.
name (str): The name of the attribute being accessed.
value (Any): The value of the attribute being accessed.
Raises:
ValueError: If the attribute is a primary key with an auto type and has not been set.
Returns:
None
"""
if name != "__fields__":
fields = object.__getattribute__(self, "__fields__")
field = fields.get(name)
if field is not None and isinstance(field.field_info, FieldInfo):
if (
field.field_info.primary_key
and field.annotation is int
and value is None
):
raise ValueError(
f"Primary key attribute '{name}' has not been set, please commit() it first."
)
def sqlmodel_validate(
cls: Type[_TSQLModel],
obj: Any,

View File

@ -79,6 +79,7 @@ from ._compat import ( # type: ignore[attr-defined]
set_config_value,
sqlmodel_init,
sqlmodel_validate,
validate_access_primary_key_autotype,
)
from .sql.sqltypes import GUID, AutoString
@ -732,6 +733,12 @@ class SQLModel(BaseModel, metaclass=SQLModelMetaclass, registry=default_registry
if name not in self.__sqlmodel_relationships__:
super().__setattr__(name, value)
def __getattribute__(self, name: str) -> Any:
# Access attributes safely using object.__getattribute__ to avoid recursion
value = object.__getattribute__(self, name)
validate_access_primary_key_autotype(self, name, value)
return value
def __repr_args__(self) -> Sequence[Tuple[Optional[str], Any]]:
# Don't show SQLAlchemy private attributes
return [