Support db apply blocking

CHANGELOG.md

@@ -110,6 +110,7 @@ Include templates data in accessible directory (failed)
 - Add ci workflow to test templates
 - Add deploy flag in model.
 - Modify the apply endpoint in the REST API to an asynchronous interface.
+- Add db apply wait on create events.
 
 #### Bug Fixes
 

superduper/backends/local/cache.py

@@ -37,11 +37,13 @@ def _put(self, component: Component):
         self._cache[component.uuid] = component
         if (component.type_id, component.identifier) in self._component_to_uuid:
             current = self._component_to_uuid[component.type_id, component.identifier]
-            current_version = self._cache[current].version
+            current_component = self._cache[current]
+            current_version = current_component.version
             if current_version < component.version:
                 self._component_to_uuid[
                     component.type_id, component.identifier
                 ] = component.uuid
+                self.expire(current_component.uuid)
         else:
             self._component_to_uuid[
                 component.type_id, component.identifier

superduper/base/datalayer.py

@@ -1,4 +1,5 @@
 import random
+import time
 import typing as t
 import warnings
 from collections import namedtuple
@@ -44,6 +45,8 @@
 PredictResult = t.Union[Document, t.Sequence[Document]]
 ExecuteResult = t.Union[SelectResult, DeleteResult, UpdateResult, InsertResult]
 
+_WAIT_TIMEOUT = 60
+
 
 class Datalayer:
     """
@@ -447,6 +450,7 @@ def apply(
         self,
         object: t.Union[Component, t.Sequence[t.Any], t.Any],
         force: bool | None = None,
+        wait: bool = True,
     ):
         """
         Add functionality in the form of components.
@@ -457,6 +461,7 @@ def apply(
         :param object: Object to be stored.
         :param dependencies: List of jobs which should execute before component
                              initialization begins.
+        :param wait: Wait for apply events.
         :return: Tuple containing the added object(s) and the original object(s).
         """
         if force is None:
@@ -537,8 +542,33 @@ def uniquify(x):
             ):
                 return object
         self.cluster.queue.publish(events=events)
+        if wait:
+            self._wait_on_events(unique_create_events)
         return object
 
+    def _wait_on_events(self, events):
+        remaining = len(events)
+        time_left = _WAIT_TIMEOUT
+        while True:
+            for event in events:
+                identifier = event.component['identifier']
+                type_id = event.component['type_id']
+                version = event.component['version']
+                try:
+                    self.show(type_id=type_id, identifier=identifier, version=version)
+                except FileNotFoundError:
+                    pass
+                else:
+                    remaining -= 1
+
+            if remaining <= 0:
+                return
+            elif time_left == 0:
+                raise TimeoutError("Timeout error while waiting for create events.")
+            else:
+                time.sleep(1)
+                time_left -= 1
+
     def remove(
         self,
         type_id: str,