Improve error messages and fix restoring optional Typed member

Author: frmdstryr

CHANGELOG.md

@@ -1,3 +1,13 @@
+# 0.8.5
+
+- Fix error that would occur when restoring optional Typed member with no value
+- Use a new `RestoreError` exception that improves error messages when the error mode is set to `raise` 
+- Fix an edge case where a cached result is not pulled when tables are joined
+
+# 0.8.4
+
+- Fix values query if it contained a related column
+
 # 0.8.3
 
 - Remove the `__ref__` member and use`id(self)` instead. State still contains the `__ref__` key.

atomdb/__init__.py

@@ -1 +1 @@
-version = "0.8.4"
+version = "0.8.5"

atomdb/base.py

@@ -497,6 +497,20 @@ def generate_getstate(cls: Type["Model"]) -> GetStateFn:
     return generate_function(source, namespace, "__getstate__")
 
 
+class RestoreError(Exception):
+    """An exception raised when an error occurs while restoring a model from state."""
+
+    def __init__(self, field: str, cls: Type["Model"], e: Exception):
+        self.field = field
+        self.cls = cls
+        self.exc = e
+
+    def __str__(self):
+        return (
+            f"Error restoring '{self.field}' on object of type {self.cls}: {self.exc}"
+        )
+
+
 def generate_restorestate(cls: Type["Model"]) -> RestoreStateFn:
     """Generate an optimized __restorestate__ function for the given model.
 
@@ -548,6 +562,7 @@ def generate_restorestate(cls: Type["Model"]) -> RestoreStateFn:
 
     namespace: DictType[str, Any] = {
         "default_unflatten": default_unflatten,
+        "RestoreError": RestoreError,
     }
     for order, f, m, unflatten in setters:
         # Since f is potentially an untrusted input, make sure it is a valid
@@ -569,7 +584,9 @@ def generate_restorestate(cls: Type["Model"]) -> RestoreStateFn:
                     RelModel = types[0]
             if RelModel is not None:
                 namespace[f"rel_model_{f}"] = RelModel
-                expr = f"await rel_model_{f}.restore(state['{f}'])"
+                expr = (
+                    f"await rel_model_{f}.restore(v) if (v := state['{f}']) else None"
+                )
             elif is_primitive_member(m):
                 # Direct assignment
                 expr = f"state['{f}']"
@@ -584,21 +601,20 @@ def generate_restorestate(cls: Type["Model"]) -> RestoreStateFn:
                 expr = f"unflatten_{f}(state['{f}'], scope)"
 
         # Do the assignment
-        if on_error == "raise":
-            template.append(f"    self.{f} = {expr}")
+        if on_error == "ignore":
+            handler = "pass"
+        elif on_error == "log":
+            handler = f"self.__log_restore_error__(e, '{f}', state, scope)"
         else:
-            if on_error == "log":
-                handler = f"self.__log_restore_error__(e, '{f}', state, scope)"
-            else:
-                handler = "pass"
-            template.extend(
-                [
-                    "    try:",
-                    f"        self.{f} = {expr}",
-                    "    except Exception as e:",
-                    f"        {handler}",
-                ]
-            )
+            handler = f"raise RestoreError('{f}', self.__class__, e) from e"
+        template.extend(
+            [
+                "    try:",
+                f"        self.{f} = {expr}",
+                "    except Exception as e:",
+                f"        {handler}",
+            ]
+        )
 
     # Update restored state
     template.append("self.__restored__ = True")
@@ -778,6 +794,7 @@ def __log_restore_error__(
     @classmethod
     async def restore(cls: Type[M], state: StateType, **kwargs: Any) -> M:
         """Restore an object from the database state"""
+        assert state is not None
         obj = cls.__new__(cls)
         await obj.__restorestate__(state)
         return obj
@@ -801,6 +818,8 @@ def flatten(self, v: Any, scope: Optional[ScopeType] = None):
         a __py__ field and arguments to reconstruct it. Also see the coercers
 
         """
+        if v is None:
+            return None
         if isinstance(v, (date, datetime, time)):
             # This is inefficient space wise but still allows queries
             s: DictType[str, Any] = {

atomdb/sql.py

@@ -60,6 +60,7 @@
     ModelManager,
     ModelMeta,
     ModelSerializer,
+    RestoreError,
     RestoreStateFn,
     ScopeType,
     StateType,
@@ -1934,11 +1935,14 @@ async def get_cached_model(cls: Type[T], pk: Any, state: StateType) -> Optional[
         If the pk is not None an instance of cls.
 
     """
-    if cls.__joined_pk__ in state and state[cls.__joined_pk__]:
+    cache = cls.objects.cache
+    if cls.__joined_pk__ in state and (joined_pk := state[cls.__joined_pk__]):
+        obj = cache.get(joined_pk)
+        if obj is not None:
+            return obj
         return await cls.restore(state)  # Restore from joined row result
     if not pk:
         return None
-    cache = cls.objects.cache
     obj = cache.get(pk)
     if obj is not None:
         return obj  # item is already in the cache
@@ -1986,6 +1990,7 @@ def generate_sql_restorestate(cls: Type["SQLModel"]) -> RestoreStateFn:
     namespace: DictType[str, Any] = {
         "default_unflatten": default_unflatten,
         "get_cached_model": get_cached_model,
+        "RestoreError": RestoreError,
     }
 
     # The state dict may have data from multiple tables that have been joined
@@ -2048,7 +2053,7 @@ def generate_sql_restorestate(cls: Type["SQLModel"]) -> RestoreStateFn:
                 # Only convert if the object has not already been restored
                 expr = "\n            ".join(
                     [
-                        f"if isinstance(v, rel_model_{f}):",
+                        f"if v is None or isinstance(v, rel_model_{f}):",
                         f"    self.{f} = v",
                         "else:",
                         f"    self.{f} = {obj}",
@@ -2067,21 +2072,20 @@ def generate_sql_restorestate(cls: Type["SQLModel"]) -> RestoreStateFn:
             else:
                 expr = f"self.{f} = unflatten_{f}(v, scope)"
 
-        if on_error == "raise":
-            template.append(f"    {expr}")
+        if on_error == "ignore":
+            handler = "pass"
+        elif on_error == "log":
+            handler = f"self.__log_restore_error__(e, '{f}', state, scope)"
         else:
-            if on_error == "log":
-                handler = f"self.__log_restore_error__(e, '{f}', state, scope)"
-            else:
-                handler = "pass"
-            template.extend(
-                [
-                    "    try:",
-                    f"        {expr}",
-                    "    except Exception as e:",
-                    f"        {handler}",
-                ]
-            )
+            handler = f"raise RestoreError('{f}', self.__class__, e) from e"
+        template.extend(
+            [
+                "    try:",
+                f"        {expr}",
+                "    except Exception as e:",
+                f"        {handler}",
+            ]
+        )
 
     # Update restored state
     template.append("self.__restored__ = True")

tests/test_base.py

@@ -24,6 +24,7 @@
     Model,
     ModelManager,
     ModelSerializer,
+    RestoreError,
     generate_function,
     is_db_field,
     is_primitive_member,
@@ -203,8 +204,10 @@ class A(Model):
         __on_error__ = "raise"
         value = Int()
 
-    with pytest.raises(TypeError):
+    with pytest.raises(RestoreError) as exc_info:
         await A.restore({"value": "str"})
+    restore_error = exc_info.value
+    assert isinstance(restore_error.exc, TypeError)
 
 
 async def test_on_error_ignore():
@@ -226,6 +229,7 @@ async def test_on_error_log(caplog):
     """
 
     class C(Model):
+        __on_error__ = "log"
         old_field = Int()
         new_field = Int()
 

tests/test_json.py

@@ -3,19 +3,25 @@
 from datetime import date, datetime, time
 from decimal import Decimal
 
+import pytest
 from atom.api import (
     Bool,
     Bytes,
     ForwardInstance,
     Instance,
+    Int,
     List,
     Range,
     Set,
     Str,
     Tuple,
+    Typed,
 )
 
-from atomdb.base import JSONModel
+from atomdb.base import JSONModel, RestoreError
+
+# Set default to raise
+JSONModel.__on_error__ = "raise"
 
 
 class Dates(JSONModel):
@@ -71,6 +77,17 @@ class Point(JSONModel):
     position = Tuple(float)
 
 
+class Position(JSONModel):
+    x = Int()
+    y = Int()
+
+
+class Rectangle(JSONModel):
+    height = Range(low=0, value=1)
+    width = Range(low=0, value=1)
+    center = Typed(Position)
+
+
 async def test_json_dates():
     now = datetime.now()
     obj = Dates(d=now.date(), t=now.time(), dt=now)
@@ -169,3 +186,31 @@ async def test_json_cyclical():
     assert r.name == "a"
     assert r.related.name == b.name
     assert r.related.related == r
+
+
+async def test_json_optional_typed_model():
+    # Test save & restore of optional typed member with no value
+    rect = Rectangle(width=1, height=2)
+    state = JSONModel.serializer.flatten(rect)
+    rect = await JSONModel.serializer.unflatten(state)
+    assert rect.center is None
+    # Now set the value and make sure it's restored'
+    rect.center = Position(x=10, y=20)
+    state = JSONModel.serializer.flatten(rect)
+    rect = await JSONModel.serializer.unflatten(state)
+    assert rect.center.x == 10 and rect.center.y == 20
+
+
+async def test_json_restore_error():
+    # Test save & restore of optional typed member with no value
+    rect = Rectangle(width=1, height=2)
+    state = JSONModel.serializer.flatten(rect)
+    state["width"] = "Not an int"
+    with pytest.raises(RestoreError) as exc_info:
+        await JSONModel.serializer.unflatten(state)
+    restore_error = exc_info.value
+    assert (
+        restore_error.field == "width"
+        and restore_error.cls is Rectangle
+        and isinstance(restore_error.exc, TypeError)
+    )

tests/test_sql.py

@@ -57,6 +57,9 @@
     SQLModelManager,
 )
 
+#: Raise errors in tests
+SQLModel.__on_error__ = "raise"
+
 
 class AbstractUser(SQLModel):
     email = Str().tag(length=64)
@@ -95,8 +98,7 @@ class JobRole(SQLModel):
     skill = Instance(JobSkill)
     tasks = Relation(lambda: JobTask)
 
-    check_one_default = sa.schema.DDL(
-        """
+    check_one_default = sa.schema.DDL("""
         CREATE OR REPLACE FUNCTION check_one_default() RETURNS TRIGGER
         LANGUAGE plpgsql
         AS $$
@@ -107,15 +109,12 @@ class JobRole(SQLModel):
             END IF;
             RETURN NEW;
         END;
-        $$;"""
-    )
+        $$;""")
 
-    trigger = sa.schema.DDL(
-        """
+    trigger = sa.schema.DDL("""
         CREATE CONSTRAINT TRIGGER check_default_role AFTER INSERT OR UPDATE
         OF "default" ON "test_sql.JobRole"
-        FOR EACH ROW EXECUTE PROCEDURE check_one_default();"""
-    )
+        FOR EACH ROW EXECUTE PROCEDURE check_one_default();""")
 
     class Meta:
         triggers = [