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,
         )
 
@@ -675,22 +678,53 @@ def _save_impl(
                 with executor:
                     for batches in executor:
                         for batch in batches:
-                            tx.send_query(batch.query, batch.args)
+                            self._send_batch_query(tx, batch)
                         batch_ids = tx.wait()
                         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, **ref.args)
+                            self._send_refetch_query(tx, ref)
 
                         refetch_data = tx.wait()
                         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_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, **ref.args)
+
     def save(
         self,
         *objs: GelModel,
@@ -718,6 +752,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
 

tests/dbsetup/chemistry.esdl

@@ -0,0 +1,70 @@
+INSERT Element {
+    name := 'Hydrogen',
+    symbol := 'H',
+    number := 1,
+    valence := 1,
+    weight := 1.0080,
+};
+INSERT Element {
+    name := 'Helium',
+    symbol := 'He',
+    number := 2,
+    valence := 0,
+    weight := 4.0026,
+};
+INSERT Element {
+    name := 'Lithium',
+    symbol := 'Li',
+    number := 3,
+    valence := 1,
+    weight := 6.94,
+};
+INSERT Element {
+    name := 'Beryllium',
+    symbol := 'Be',
+    number := 4,
+    valence := 2,
+    weight := 9.0122,
+};
+INSERT Element {
+    name := 'Boron',
+    symbol := 'B',
+    number := 5,
+    valence := 3,
+    weight := 10.81,
+};
+INSERT Element {
+    name := 'Carbon',
+    symbol := 'C',
+    number := 6,
+    valence := 4,
+    weight := 12.011,
+};
+INSERT Element {
+    name := 'Nitrogen',
+    symbol := 'N',
+    number := 7,
+    valence := 5,
+    weight := 14.007,
+};
+INSERT Element {
+    name := 'Oxygen',
+    symbol := 'O',
+    number := 8,
+    valence := 6,
+    weight := 15.999,
+};
+INSERT Element {
+    name := 'Fluorine',
+    symbol := 'F',
+    number := 9,
+    valence := 7,
+    weight := 18.998,
+};
+INSERT Element {
+    name := 'Neon',
+    symbol := 'Ne',
+    number := 10,
+    valence := 0,
+    weight := 20.180,
+};

tests/dbsetup/chemistry.gel

@@ -0,0 +1,63 @@
+module default {
+    type Named {
+        required name: std::str;
+    };
+
+    type Element extending Named {
+        required symbol: str;
+        required number: int64;
+        required valence: int64;
+        required weight: float64;
+    };
+
+    # We do this for now, since self multi links aren't well supported
+    type BaseAtom {
+        # Single link no props
+        required element: Element;
+
+        # Computed single prop
+        weight := sum(.element.weight);
+    };
+
+    type RefAtom extending BaseAtom {
+        # Multi link with props
+        multi bonds: BaseAtom {
+            count: int64;
+        }
+
+        # Computed single link
+        compound := .<atoms[is Compound];
+    };
+
+    type Compound extending Named {
+        # Multi link no props
+        required multi atoms: RefAtom;
+
+        # Multi prop
+        multi alternate_names: str;
+    };
+
+    type Reactor {
+        # Computed multi link
+        multi atoms := .<reactor[is Atom];
+
+        # Computed single prop
+        total_weight := sum(.atoms.element.weight);
+        # Computed multi prop
+        atom_weights := .atoms.element.weight;
+    };
+
+    type Atom extending BaseAtom {
+        # Multi link with props
+        multi bonds: BaseAtom {
+            count: int64;
+        }
+
+        # Single link with props
+        required reactor: Reactor;
+
+        # Computed single prop using link props
+        total_bond_count := sum(.bonds@count);
+        total_bond_weight := sum(.bonds.weight);
+    };
+}