From e36268b9c69f9d495a26fb2b9af711a8bc131358 Mon Sep 17 00:00:00 2001 From: JieguangZhou Date: Tue, 2 Jan 2024 14:50:31 +0800 Subject: [PATCH] Optimization function add_component_to_cache --- superduperdb/base/datalayer.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/superduperdb/base/datalayer.py b/superduperdb/base/datalayer.py index a5eefbe0dc..3298a1cc3f 100644 --- a/superduperdb/base/datalayer.py +++ b/superduperdb/base/datalayer.py @@ -912,7 +912,7 @@ def _add( if parent is not None: self.metadata.create_parent_child(parent, object.unique_id) object.post_create(self) - self.add_component_to_cache(object) + self._add_component_to_cache(object) these_jobs = object.schedule_jobs(self, dependencies=dependencies) jobs.extend(these_jobs) return jobs @@ -1263,15 +1263,15 @@ def close(self): # TODO: gracefully close all opened connections return - def add_component_to_cache(self, component: Component): + def _add_component_to_cache(self, component: Component): + """ + Add component to cache when it is added to the db. + Avoiding the need to load it from the db again. + """ type_id = component.type_id - - # TODO: Only added model, - # if add other component will get an error, need to fix - if type_id != "model": - return - - self.models[component.identifier] = component + if cm := self.type_id_to_cache_mapping.get(type_id): + getattr(self, cm)[component.identifier] = component + component.on_load(self) @dc.dataclass