From 435f699fb1113ab259ef7970ecbaae849e38b71c Mon Sep 17 00:00:00 2001
From: tianwei <liutianweidlut@gmail.com>
Date: Tue, 18 Jul 2023 11:07:36 +0800
Subject: [PATCH] enhance(sdk): the localstore data will be auto dumpped when
 close the tablewriter (#2511)

---
 client/starwhale/api/_impl/data_store.py |  7 ++++++-
 client/tests/core/test_dataset.py        | 10 ++++++++--
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/client/starwhale/api/_impl/data_store.py b/client/starwhale/api/_impl/data_store.py
index 64474013f7..d96ac9f001 100644
--- a/client/starwhale/api/_impl/data_store.py
+++ b/client/starwhale/api/_impl/data_store.py
@@ -1217,6 +1217,7 @@ def loads(cls, file: Path, table_name: str) -> MemoryTable:
 
     def dump(self, root_path: str, if_dirty: bool = True) -> None:
         root = Path(root_path)
+        ensure_dir(root)
         lock = root / ".lock"
         with FileLock(lock):
             return self._dump(root, if_dirty)
@@ -1696,7 +1697,11 @@ def close(self) -> None:
                 self._stopped = True
                 self._cond.notify()
         self.join()
-        self._raise_run_exceptions(0)
+        try:
+            if isinstance(self.data_store, LocalDataStore):
+                self.data_store.dump()
+        finally:
+            self._raise_run_exceptions(0)
 
     def _raise_run_exceptions(self, limits: int) -> None:
         if len(self._queue_run_exceptions) > limits:
diff --git a/client/tests/core/test_dataset.py b/client/tests/core/test_dataset.py
index 6e75660515..2ad58d7982 100644
--- a/client/tests/core/test_dataset.py
+++ b/client/tests/core/test_dataset.py
@@ -121,7 +121,8 @@ def test_build_from_yaml(self, m_import: MagicMock) -> None:
         assert call_args[1].name == "mnist"
         assert m_import.call_args[0][1] == "dataset:build"
 
-    def test_build_from_image_folder(self) -> None:
+    @patch("starwhale.api._impl.data_store.LocalDataStore.dump")
+    def test_build_from_image_folder(self, m_dump: MagicMock) -> None:
         image_folder = Path("/tmp/workdir/images")
         ensure_file(image_folder / "1.jpg", "1", parents=True)
         ensure_file(image_folder / "1.txt", "1", parents=True)
@@ -159,6 +160,7 @@ def test_build_from_image_folder(self) -> None:
             alignment_size="128",
             volume_size="128M",
         )
+        assert m_dump.call_count == 1
 
     @patch("starwhale.api._impl.dataset.model.Dataset.from_huggingface")
     def test_build_from_huggingface(self, m_hf: MagicMock) -> None:
@@ -314,6 +316,7 @@ def test_build_from_handler(self, m_import: MagicMock) -> None:
         assert call_args[1].attr.volume_size == D_FILE_VOLUME_SIZE
 
     @patch("starwhale.base.uri.resource.Resource._refine_local_rc_info")
+    @patch("starwhale.api._impl.data_store.LocalDataStore.dump")
     def test_build_workflow(self, *args: t.Any) -> None:
         class _MockBuildExecutor:
             def __iter__(self) -> t.Generator:
@@ -399,6 +402,7 @@ def __iter__(self) -> t.Generator:
         assert len(os.listdir(sw.rootdir / SW_TMP_DIR_NAME)) == 0
 
     @patch("starwhale.base.uri.resource.Resource._refine_local_rc_info")
+    @patch("starwhale.api._impl.data_store.LocalDataStore.dump")
     def test_head(self, *args: t.Any) -> None:
         from starwhale.api._impl.dataset import Dataset as SDKDataset
 
@@ -447,7 +451,8 @@ def test_head(self, *args: t.Any) -> None:
         DatasetTermViewJson(dataset_uri).head(2, show_raw_data=True)
 
     @patch("starwhale.base.uri.resource.Resource._refine_local_rc_info")
-    def test_from_json(self, *args: t.Any) -> None:
+    @patch("starwhale.api._impl.data_store.LocalDataStore.dump")
+    def test_from_json(self, m_dump: MagicMock, *args: t.Any) -> None:
         from starwhale.api._impl.dataset import Dataset as SDKDataset
 
         myds = SDKDataset.from_json(
@@ -462,6 +467,7 @@ def test_from_json(self, *args: t.Any) -> None:
             "content.child_content",
         )
         assert myds[1].features["zh-cn"] == "最近怎么样"
+        assert m_dump.call_count == 2
 
 
 class TestJsonDict(TestCase):