Add profiling for sync.

Author: dnwpark

gel/_internal/_save.py

@@ -1074,6 +1074,7 @@ def make_save_executor_constructor(
     refetch: bool,
     warn_on_large_sync_set: bool = False,
     save_postcheck: bool = False,
+    executor_type: type,
 ) -> Callable[[], SaveExecutor]:
     (
         create_batches,
@@ -1085,7 +1086,7 @@ def make_save_executor_constructor(
         refetch=refetch,
         warn_on_large_sync_set=warn_on_large_sync_set,
     )
-    return lambda: SaveExecutor(
+    return lambda: executor_type(
         objs=objs,
         create_batches=create_batches,
         updates=updates,

gel/_testbase.py

@@ -479,17 +479,21 @@ def make_test_client(
                 if cls.is_client_async
                 else blocking_client.BlockingIOConnection
             )
-        client_class = (
-            TestAsyncIOClient
-            if issubclass(connection_class, asyncio_client.AsyncIOConnection)
-            else TestClient
-        )
+        client_class = cls._get_client_class(connection_class)
         return client_class(
             connection_class=connection_class,
             max_concurrency=1,
             **conargs,
         )
 
+    @classmethod
+    def _get_client_class(cls, connection_class):
+        return (
+            TestAsyncIOClient
+            if issubclass(connection_class, asyncio_client.AsyncIOConnection)
+            else TestClient
+        )
+
     @classmethod
     def get_connect_args(
         cls, *, cluster=None, database="edgedb", user="edgedb", password="test"

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
@@ -600,6 +600,7 @@ class AsyncIOClient(
 
     __slots__ = ()
     _impl_class = _AsyncIOPoolImpl
+    _save_executor_type = SaveExecutor
 
     async def check_connection(self) -> base_client.ConnectionInfo:
         return await self._impl.ensure_connected()
@@ -647,6 +648,7 @@ async def _save_impl(
             refetch=refetch,
             save_postcheck=opts.save_postcheck,
             warn_on_large_sync_set=warn_on_large_sync_set,
+            executor_type=self._save_executor_type,
         )
 
         async for tx in self._batch():

gel/blocking_client.py

@@ -39,7 +39,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
@@ -651,6 +656,7 @@ class Client(
 
     __slots__ = ()
     _impl_class = _PoolImpl
+    _save_executor_type = SaveExecutor
 
     def _save_impl(
         self,
@@ -659,12 +665,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,
         )
 
@@ -674,28 +677,74 @@ def _save_impl(
 
                 with 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.feed_db_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
                         ):
                             ref.feed_db_data(ref_data)
 
+    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,
@@ -723,6 +772,7 @@ def __debug_save__(self, *objs: GelModel) -> SaveDebug:
         make_executor = make_save_executor_constructor(
             objs,
             refetch=False,  # TODO
+            executor_type=self._save_executor_type,
         )
         plan_time = time.monotonic_ns() - ns