Add profiling for sync.

dnwpark · dnwpark · commit c8a2b763e6b3 · 2025-09-17T15:40:31.000-07:00
diff --git a/gel/_internal/_save.py b/gel/_internal/_save.py
@@ -1160,13 +1160,14 @@ def make_save_executor_constructor(
     refetch: bool,
     warn_on_large_sync_set: bool = False,
     save_postcheck: bool = False,
+    executor_type: type,
 ) -> Callable[[], SaveExecutor]:
     plan = make_plan(
         objs,
         refetch=refetch,
         warn_on_large_sync_set=warn_on_large_sync_set,
     )
-    return lambda: SaveExecutor(
+    return lambda: executor_type(
         objs=objs,
         create_batches=plan.create_batches,
         updates=plan.update_batch,
diff --git a/gel/_internal/_testbase/_base.py b/gel/_internal/_testbase/_base.py
@@ -630,6 +630,7 @@ class BranchTestCase(InstanceTestCase):
     BASE_TEST_CLASS = True
     TEARDOWN_RETRY_DROP_DB = 1
 
+    CLIENT_TYPE: ClassVar[type[TestClient | TestAsyncIOClient] | None]
     client: ClassVar[TestClient | TestAsyncIOClient]
 
     @classmethod
@@ -678,7 +679,9 @@ def setUp(self) -> None:
         if self.ISOLATED_TEST_BRANCHES:
             cls = type(self)
             testdb = cls.loop.run_until_complete(self.setup_branch_copy())
-            client = cls.make_test_client(database=testdb)._with_debug(
+            client = cls.make_test_client(
+                database=testdb, client_class=self.CLIENT_TYPE
+            )._with_debug(
                 save_postcheck=True,
             )
             self.client = client  # type: ignore[misc]
@@ -717,6 +720,7 @@ def tearDown(self) -> None:
     def make_test_client(
         cls,
         *,
+        client_class: type[TestClient | TestAsyncIOClient] | None = None,
         connection_class: type[
             asyncio_client.AsyncIOConnection
             | blocking_client.BlockingIOConnection
@@ -758,14 +762,17 @@ def make_blocking_test_client(
         cls,
         *,
         instance: _server.BaseInstance,
+        client_class: type[TestClient] | None = None,
         connection_class: type[blocking_client.BlockingIOConnection]
         | None = None,
         **kwargs: str,
     ) -> TestClient:
+        if client_class is None:
+            client_class = TestClient
         if connection_class is None:
             connection_class = blocking_client.BlockingIOConnection
         client = instance.create_blocking_client(
-            client_class=TestClient,
+            client_class=client_class,
             connection_class=connection_class,
             **cls.get_connect_args(instance, **kwargs),
         )
@@ -799,13 +806,16 @@ def make_async_test_client(
         cls,
         *,
         instance: _server.BaseInstance,
+        client_class: type[TestAsyncIOClient] | None = None,
         connection_class: type[asyncio_client.AsyncIOConnection] | None = None,
         **kwargs: str,
     ) -> TestAsyncIOClient:
+        if client_class is None:
+            client_class = TestAsyncIOClient
         if connection_class is None:
             connection_class = asyncio_client.AsyncIOConnection
         client = instance.create_async_client(
-            client_class=TestAsyncIOClient,
+            client_class=client_class,
             connection_class=connection_class,
             **cls.get_connect_args(instance, **kwargs),
         )
@@ -881,7 +891,9 @@ async def setup_and_connect(cls) -> None:
                 await cls._create_empty_branch(dbname)
 
         if not cls.ISOLATED_TEST_BRANCHES:
-            cls.client = cls.make_test_client(database=dbname)
+            cls.client = cls.make_test_client(
+                database=dbname, client_class=cls.CLIENT_TYPE
+            )
             if isinstance(cls.client, gel.AsyncIOClient):
                 await cls.client.ensure_connected()
             else:
@@ -1021,11 +1033,13 @@ class AsyncQueryTestCase(BranchTestCase):
     def make_test_client(  # pyright: ignore [reportIncompatibleMethodOverride]
         cls,
         *,
+        client_class: type[TestAsyncIOClient] | None = None,
         connection_class: type[asyncio_client.AsyncIOConnection] | None = None,  # type: ignore [override]
         **kwargs: str,
     ) -> TestAsyncIOClient:
         return cls.make_async_test_client(
             instance=cls.instance,
+            client_class=client_class,
             connection_class=connection_class,
             **kwargs,
         )
@@ -1062,12 +1076,14 @@ def adapt_call(cls, coro: Any) -> Any:
     def make_test_client(  # pyright: ignore [reportIncompatibleMethodOverride]
         cls,
         *,
+        client_class: type[TestClient] | None = None,
         connection_class: type[blocking_client.BlockingIOConnection]  # type: ignore [override]
         | None = None,
         **kwargs: str,
     ) -> TestClient:
         return cls.make_blocking_test_client(
             instance=cls.instance,
+            client_class=client_class,
             connection_class=connection_class,
             **kwargs,
         )
diff --git a/gel/asyncio_client.py b/gel/asyncio_client.py
@@ -37,7 +37,7 @@
 from .protocol import asyncio_proto  # type: ignore [attr-defined, unused-ignore]
 from .protocol.protocol import InputLanguage, OutputFormat
 
-from ._internal._save import make_save_executor_constructor
+from ._internal._save import make_save_executor_constructor, SaveExecutor
 
 if typing.TYPE_CHECKING:
     from ._internal._qbmodel._pydantic import GelModel
@@ -675,6 +675,7 @@ async def _save_impl(
             refetch=refetch,
             save_postcheck=opts.save_postcheck,
             warn_on_large_sync_set=warn_on_large_sync_set,
+            executor_type=SaveExecutor,
         )
 
         async for tx in self._batch():
diff --git a/gel/blocking_client.py b/gel/blocking_client.py
@@ -38,7 +38,12 @@
 from .protocol import blocking_proto  # type: ignore [attr-defined, unused-ignore]
 from .protocol.protocol import InputLanguage, OutputFormat
 
-from ._internal._save import make_save_executor_constructor
+from ._internal._save import (
+    QueryBatch,
+    QueryRefetch,
+    SaveExecutor,
+    make_save_executor_constructor,
+)
 
 if typing.TYPE_CHECKING:
     from ._internal._qbmodel._pydantic import GelModel
@@ -681,6 +686,7 @@ class Client(
 
     __slots__ = ()
     _impl_class = _PoolImpl
+    _save_executor_type = SaveExecutor
 
     def _save_impl(
         self,
@@ -689,12 +695,9 @@ def _save_impl(
         objs: tuple[GelModel, ...],
         warn_on_large_sync_set: bool = False,
     ) -> None:
-        opts = self._get_debug_options()
-
-        make_executor = make_save_executor_constructor(
-            objs,
+        make_executor = self._get_make_save_executor(
             refetch=refetch,
-            save_postcheck=opts.save_postcheck,
+            objs=objs,
             warn_on_large_sync_set=warn_on_large_sync_set,
         )
 
@@ -703,23 +706,13 @@ def _save_impl(
                 executor = make_executor()
 
                 for batches in executor:
-                    for batch in batches:
-                        tx.send_query(batch.query, batch.args)
-                    batch_ids = tx.wait()
+                    batch_ids = self._send_batch_queries(tx, batches)
                     for ids, batch in zip(batch_ids, batches, strict=True):
                         batch.record_inserted_data(ids)
 
                 if refetch:
                     ref_queries = executor.get_refetch_queries()
-                    for ref in ref_queries:
-                        tx.send_query(
-                            ref.query,
-                            spec=ref.args.spec,
-                            new=ref.args.new,
-                            existing=ref.args.existing,
-                        )
-
-                    refetch_data = tx.wait()
+                    refetch_data = self._send_refetch_queries(tx, ref_queries)
 
                     for ref_data, ref in zip(
                         refetch_data, ref_queries, strict=True
@@ -728,6 +721,60 @@ def _save_impl(
 
         executor.commit()
 
+    def _get_make_save_executor(
+        self,
+        *,
+        refetch: bool,
+        objs: tuple[GelModel, ...],
+        warn_on_large_sync_set: bool = False,
+    ) -> typing.Callable[[], SaveExecutor]:
+        opts = self._get_debug_options()
+
+        return make_save_executor_constructor(
+            objs,
+            refetch=refetch,
+            save_postcheck=opts.save_postcheck,
+            warn_on_large_sync_set=warn_on_large_sync_set,
+            executor_type=self._save_executor_type,
+        )
+
+    def _send_batch_queries(
+        self,
+        tx: BatchIteration,
+        batches: list[QueryBatch],
+    ) -> list[Any]:
+        for batch in batches:
+            self._send_batch_query(tx, batch)
+        return tx.wait()
+
+    def _send_refetch_queries(
+        self,
+        tx: BatchIteration,
+        ref_queries: list[QueryRefetch],
+    ) -> list[Any]:
+        for ref in ref_queries:
+            self._send_refetch_query(tx, ref)
+        return tx.wait()
+
+    def _send_batch_query(
+        self,
+        tx: BatchIteration,
+        batch: QueryBatch,
+    ) -> None:
+        tx.send_query(batch.query, batch.args)
+
+    def _send_refetch_query(
+        self,
+        tx: BatchIteration,
+        ref: QueryRefetch,
+    ) -> None:
+        tx.send_query(
+            ref.query,
+            spec=ref.args.spec,
+            new=ref.args.new,
+            existing=ref.args.existing,
+        )
+
     def save(
         self,
         *objs: GelModel,
diff --git a/tests/test_profiling.py b/tests/test_profiling.py
