Add QueryRefetchArgs. (#863)

A simple refactor for _save.

Author: dnwpark

gel/_internal/_save.py

@@ -1156,11 +1156,40 @@ def feed_db_data(self, obj_data: Iterable[Any]) -> None:
                 pass
 
 
+# Refetching links will only refetch linked objects which are already in the
+# link, or are in the sync objects.
+#
+# A refetched object will therefore need, per link:
+# - whether that link is being being refetched, and if so,
+# - the objects already in the link.
+#
+# The order of links is kept consistent with the refetch shape.
+RefetchSpecEntry = TypeAliasType(
+    "RefetchSpecEntry",
+    tuple[
+        uuid.UUID,  # Refetched obj id
+        list[
+            tuple[
+                bool,  # Whether the link is being refetched
+                list[uuid.UUID],  # Objects in the link
+            ]
+        ],
+    ],
+)
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class QueryRefetchArgs:
+    spec: list[RefetchSpecEntry]
+    new: list[uuid.UUID]
+    existing: list[uuid.UUID]
+
+
 @_struct
 class QueryRefetch:
     executor: SaveExecutor
     query: TypeWrapper[type[GelModel]]
-    args: dict[str, Any]
+    args: QueryRefetchArgs
     shape: RefetchShape
 
     def feed_db_data(self, obj_data: Iterable[Any]) -> None:
@@ -1305,7 +1334,7 @@ def _compile_refetch(
         tuple[
             type[GelModel],
             TypeWrapper[type[GelModel]],
-            list[Any],
+            list[RefetchSpecEntry],
             RefetchShape,
         ]
     ]:
@@ -1411,7 +1440,7 @@ def _compile_refetch(
                         else:
                             obj_link_ids.append(self._get_id(s_link))
 
-            spec_arg = [
+            spec_arg: list[RefetchSpecEntry] = [
                 (
                     obj_id,
                     [
@@ -1469,11 +1498,11 @@ def get_refetch_queries(
             QueryRefetch(
                 executor=self,
                 query=q[1],
-                args={
-                    "spec": q[2],
-                    "new": new_ids,
-                    "existing": list(self.existing_objects),
-                },
+                args=QueryRefetchArgs(
+                    spec=q[2],
+                    new=new_ids,
+                    existing=list(self.existing_objects),
+                ),
                 shape=q[3],
             )
             for q in self._compile_refetch()

gel/asyncio_client.py

@@ -664,7 +664,12 @@ async def _save_impl(
                     if refetch:
                         ref_queries = executor.get_refetch_queries()
                         for ref in ref_queries:
-                            await tx.send_query(ref.query, **ref.args)
+                            await tx.send_query(
+                                ref.query,
+                                spec=ref.args.spec,
+                                new=ref.args.new,
+                                existing=ref.args.existing,
+                            )
 
                         refetch_data = await tx.wait()
                         for ref_data, ref in zip(

gel/blocking_client.py

@@ -683,7 +683,12 @@ def _save_impl(
                     if refetch:
                         ref_queries = executor.get_refetch_queries()
                         for ref in ref_queries:
-                            tx.send_query(ref.query, **ref.args)
+                            tx.send_query(
+                                ref.query,
+                                spec=ref.args.spec,
+                                new=ref.args.new,
+                                existing=ref.args.existing,
+                            )
 
                         refetch_data = tx.wait()
                         for ref_data, ref in zip(