]> git.ipfire.org Git - thirdparty/fastapi/sqlmodel.git/commitdiff
✨ Add support for safe access to PK int autotypes
authorEsteban Maya Cadavid <emayacadavid9@gmail.com>
Mon, 1 Jul 2024 22:28:11 +0000 (17:28 -0500)
committerEsteban Maya Cadavid <emayacadavid9@gmail.com>
Mon, 1 Jul 2024 22:28:11 +0000 (17:28 -0500)
sqlmodel/_compat.py
sqlmodel/main.py

index 4018d1bb396eb1d963b6951ed20317f09c1b6dcc..02d56135eea0cbcdedb207849971a21ce502662c 100644 (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,
index 505683f756c718783943ee6deaea7e59546748a3..6a5a46a67c9b3802af18c92a76f6a4de0dc4e382 100644 (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 [