Refetch single links when syncing (#911)

If a single link is changed as a result of a sync insert or update, the resulting link was always None.

If the refetched link target is not passed as a sync argument, the refetch will now put an object with only `id` set.

Author: dnwpark

gel/_internal/_qbmodel/_abstract/_descriptors.py

@@ -771,7 +771,11 @@ def reconcile_link(
         else:
             return next(iter(link_to))
 
-    return new_objects[obj_id]
+    if obj_id in new_objects:
+        return new_objects[obj_id]
+
+    # The refetched object appeared out of nowhere.
+    return refetched
 
 
 def reconcile_proxy_link(
@@ -799,8 +803,11 @@ def reconcile_proxy_link(
             link_to = _link_to
         else:
             link_to = next(iter(_link_to))
-    else:
+    elif obj_id in new_objects:
         link_to = new_objects[obj_id]
+    else:
+        # The refetched object appeared out of nowhere.
+        link_to = refetched.without_linkprops()
 
     # Make sure we create a new proxy model that
     # would wrap either an object that already exists

gel/_internal/_save.py

@@ -640,12 +640,18 @@ def _add_refetch_shape(
             ref_shape.fields[field_name] = ptr_info
 
     if obj.__gel_new__:
-        # New objects should additionally refetch any computed properties
         for ptr_name, ptr_info in tp_pointers.items():
-            if not ptr_info.computed or ptr_info.kind != PointerKind.Property:
-                continue
+            # New objects should refetch computed properties
+            if ptr_info.kind == PointerKind.Property and ptr_info.computed:
+                ref_shape.fields[ptr_name] = ptr_info
 
-            ref_shape.fields[ptr_name] = ptr_info
+            # New objects should refetch single links
+            if (
+                ptr_info.kind == PointerKind.Link
+                and not ptr_info.cardinality.is_multi()
+                and ptr_name != "__type__"
+            ):
+                ref_shape.fields[ptr_name] = ptr_info
 
 
 def push_refetch_new(
@@ -1347,13 +1353,17 @@ def _compile_refetch(
                     link_arg_order.append(ptr.name)
                     link_num = len(link_arg_order) - 1
 
+                    link_shape_elems: list[str] = []
+                    if not ptr.cardinality.is_multi():
+                        link_shape_elems += ["id"]
                     if ptr.properties:
-                        props = ",".join(
+                        link_shape_elems += [
                             f"@{quote_ident(p)}" for p in ptr.properties
-                        )
-                        props = f"{{ {props} }}"
-                    else:
-                        props = ""
+                        ]
+
+                    link_shape = ""
+                    if link_shape_elems:
+                        link_shape = f"{{ {','.join(link_shape_elems)} }}"
 
                     maybe_assert = maybe_assert_end = ""
 
@@ -1365,25 +1375,34 @@ def _compile_refetch(
                     if ptr.computed:
                         computed_filter = "or .id in __existing"
 
-                    # The logic of the `filter` clause is:
-                    # - first test if the link was fetched for this object
-                    #   at all - if not - we don't need to update it
-                    # - if the link was fetched for this particular GelModel --
-                    #   filter it by all existing objects prior to sync() in
-                    #   this link PLUS all new objects that were inserted
-                    #   during sync().
-                    select_shape.append(f"""
-                        {quote_ident(ptr.name)} := {maybe_assert}(
-                            with
-                                __link := __obj_data.1[{link_num}]
-                            select .{quote_ident(ptr.name)} {props}
-                            filter __link.0 and (
-                                .id in array_unpack(__link.1)
-                                or .id in __new
-                                {computed_filter}
-                            )
-                        ){maybe_assert_end}
-                    """)
+                    if ptr.cardinality.is_multi():
+                        # The logic of the `filter` clause is:
+                        # - first test if the link was fetched for this object
+                        #   at all - if not - we don't need to update it
+                        # - if the link was fetched for this particular
+                        #   GelModel -- filter it by all existing objects prior
+                        #   to sync() in this link PLUS all new objects that
+                        #   were inserted during sync().
+                        select_shape.append(f"""
+                            {quote_ident(ptr.name)} := {maybe_assert}(
+                                with
+                                    __link := __obj_data.1[{link_num}]
+                                select .{quote_ident(ptr.name)} {link_shape}
+                                filter __link.0 and (
+                                    .id in array_unpack(__link.1)
+                                    or .id in __new
+                                    {computed_filter}
+                                )
+                            ){maybe_assert_end}
+                        """)
+
+                    else:
+                        # Single links should always refetch the link
+                        select_shape.append(f"""
+                            {quote_ident(ptr.name)} := {maybe_assert}(
+                                .{quote_ident(ptr.name)} {link_shape}
+                            ){maybe_assert_end}
+                        """)
 
                 else:
                     select_shape.append(quote_ident(ptr.name))
@@ -1568,6 +1587,8 @@ def _apply_refetched_data_shape(
 
             model_or_models: GelModel | IDTracker[GelModel, None]
             if gel_id in self.new_objects:
+                # For new objects, update the data we got back from the batch
+                # queries.
                 model_or_models = self.new_objects[gel_id]
             else:
                 model_or_models = self.existing_objects[gel_id]
@@ -1593,66 +1614,86 @@ def _apply_refetched_data_shape(
                     continue
 
                 shape_field = shape.fields[field]
-                is_multi = shape_field.cardinality.is_multi()
-                is_link = shape_field.kind.is_link()
 
                 for model in models:
                     assert model.id == gel_id
-                    model_dict = model.__dict__
-
-                    model.__pydantic_fields_set__.add(field)
-
-                    if is_link and is_multi:
-                        link = model_dict.get(field)
-                        if link is None:
-                            # This instance never had this link, but
-                            # now it will.
-                            # TODO: this needs to be optimized.
-                            link = copy.copy(new_value)
-                            model_dict[field] = link
-                            link.__gel_replace_with_empty__()
-
-                        assert is_link_abstract_dlist(link)
-                        assert type(new_value) is type(link)
-                        model_dict[field] = link.__gel_reconcile__(
-                            # no need to copy `new_value`,
-                            # it will only be iterated over
-                            new_value,
-                            self.existing_objects,  # type: ignore [arg-type]
-                            self.new_objects,  # type: ignore [arg-type]
-                        )
-
-                    elif is_link:
-                        if new_value is None:
-                            model_dict[field] = new_value
-                        elif shape_field.properties:
-                            model_dict[field] = reconcile_proxy_link(
-                                existing=model_dict.get(field),
-                                refetched=new_value,  # pyright: ignore [reportArgumentType]
-                                existing_objects=self.existing_objects,  # type: ignore [arg-type]
-                                new_objects=self.new_objects,
-                            )
-                        else:
-                            model_dict[field] = reconcile_link(
-                                existing=model_dict.get(field),
-                                refetched=new_value,
-                                existing_objects=self.existing_objects,  # type: ignore [arg-type]
-                                new_objects=self.new_objects,
-                            )
 
-                    elif not is_link:
-                        # could be a multi prop, could be a single prop
-                        # with an array value
-                        if shape_field.mutable:
-                            model_dict[field] = deepcopy(new_value)
-                        else:
-                            model_dict[field] = new_value
+                    self._apply_refetched_field_to_model(
+                        shape_field, new_value, model, deepcopy
+                    )
 
             # Let's be extra cautious and help GC a bit here. `obj` can
             # have recursive references to other objects via links and
             # computed backlinks.
             obj.__dict__.clear()
 
+    def _apply_refetched_field_to_model(
+        self,
+        shape_field: GelPointerReflection,
+        new_value: Any,
+        model: GelModel,
+        deepcopy: Callable[[T], T],
+    ) -> None:
+        field = shape_field.name
+        is_multi = shape_field.cardinality.is_multi()
+        is_link = shape_field.kind.is_link()
+
+        model_dict = model.__dict__
+
+        model.__pydantic_fields_set__.add(field)
+
+        if is_link and is_multi:
+            link = model_dict.get(field)
+            assert isinstance(new_value, AbstractCollection)
+            if link is None:
+                # This instance never had this link, but
+                # now it will.
+                # TODO: this needs to be optimized.
+                link = copy.copy(new_value)
+                model_dict[field] = link
+                link.__gel_replace_with_empty__()
+
+            assert is_link_abstract_dlist(link)
+            assert type(new_value) is type(link)
+            model_dict[field] = link.__gel_reconcile__(
+                # no need to copy `new_value`,
+                # it will only be iterated over
+                new_value,  # type: ignore [arg-type]
+                self.existing_objects,  # type: ignore [arg-type]
+                self.new_objects,  # type: ignore [arg-type]
+            )
+
+        elif is_link:
+            existing = model_dict.get(field)
+            if (
+                new_value is None
+                or existing is None
+                or existing is DEFAULT_VALUE
+            ):
+                model_dict[field] = new_value
+            elif shape_field.properties:
+                model_dict[field] = reconcile_proxy_link(
+                    existing=existing,
+                    refetched=new_value,  # pyright: ignore [reportArgumentType]
+                    existing_objects=self.existing_objects,  # type: ignore [arg-type]
+                    new_objects=self.new_objects,
+                )
+            else:
+                model_dict[field] = reconcile_link(
+                    existing=existing,
+                    refetched=new_value,
+                    existing_objects=self.existing_objects,  # type: ignore [arg-type]
+                    new_objects=self.new_objects,
+                )
+
+        elif not is_link:
+            # could be a multi prop, could be a single prop
+            # with an array value
+            if shape_field.mutable:
+                model_dict[field] = deepcopy(new_value)
+            else:
+                model_dict[field] = new_value
+
     def _commit(self) -> None:
         for obj, new_id in self.new_object_ids.items():
             assert new_id is not None
@@ -1662,25 +1703,36 @@ def _commit(self) -> None:
         if self.refetch:
             self._apply_refetched_data()
 
+            # Apply refetched data to new objects.
+            #
+            # There are 3 instances of a newly synced object:
+            # - The user instance
+            # - The batch instance
+            # - The refetch instance
+            #
+            # In _apply_refetched_data, the batch instance will be updated
+            # using data from the refetech instance. Here, we finally update
+            # the user's copy.
             for obj, new_id in self.new_object_ids.items():
                 assert new_id is not None
 
                 updated = self.new_objects[new_id]
-                pydantic_set_fields = obj.__pydantic_fields_set__
 
                 for prop in get_pointers(type(obj)):
                     if (
                         # id is already set and should never change.
                         prop.name == "id"
                         # prop not refetched
                         or prop.name not in updated.__dict__
-                        # TODO: Refetching links for new objects
-                        or prop.kind == PointerKind.Link
                     ):
                         continue
 
-                    obj.__dict__[prop.name] = updated.__dict__[prop.name]
-                    pydantic_set_fields.add(prop.name)
+                    self._apply_refetched_field_to_model(
+                        prop,
+                        updated.__dict__[prop.name],
+                        obj,
+                        _identity_func,
+                    )
 
         self._commit_recursive()
 

tests/test_model_generator.py

@@ -3916,26 +3916,18 @@ def test_modelgen_save_reload_links_04(self):
 
         self.client.sync(tsl)
 
-        # Computed links should not be set -- we never fetched them
-        with self.assertRaisesRegex(
-            AttributeError, "'comp_req_wprop_friend' is not set"
-        ):
-            assert tsl.comp_req_wprop_friend is not None  # access field
+        # Computed single links should be refetched, but the ids
+        self.assertEqual(tsl.comp_req_wprop_friend, alice)
+        self.assertFalse(hasattr(tsl.comp_req_wprop_friend, "name"))
 
-        with self.assertRaisesRegex(
-            AttributeError, "'comp_req_friend' is not set"
-        ):
-            assert tsl.comp_req_friend is not None  # access field
+        self.assertEqual(tsl.comp_req_friend, alice)
+        self.assertFalse(hasattr(tsl.comp_req_friend, "name"))
 
-        with self.assertRaisesRegex(
-            AttributeError, "'comp_opt_friend' is not set"
-        ):
-            assert tsl.comp_opt_friend is not None  # access field
+        self.assertEqual(tsl.comp_opt_friend, alice)
+        self.assertFalse(hasattr(tsl.comp_opt_friend, "name"))
 
-        with self.assertRaisesRegex(
-            AttributeError, "'comp_opt_wprop_friend' is not set"
-        ):
-            assert tsl.comp_opt_wprop_friend is not None  # access field
+        self.assertEqual(tsl.comp_opt_wprop_friend, alice)
+        self.assertFalse(hasattr(tsl.comp_opt_wprop_friend, "name"))
 
         # Make mypy happy
         assert tsl.opt_friend is not None
@@ -4203,8 +4195,6 @@ def test_modelgen_save_reload_links_08(self):
         )
         self.assertEqual({gr.id for gr in alice.groups}, orig_groups | {g.id})
 
-    @tb.xfail
-    # link props aren't reloaded after save
     def test_modelgen_save_reload_linkprops_01(self):
         from models.orm import default