fix(storage): preserve tree on invalid or failed moves

chat2db-community-server/chat2db-community-storage/src/main/java/ai/chat2db/community/storage/small/DataSourceStorage.java

@@ -238,7 +238,7 @@ protected synchronized void createDataSourceNode(Long datasourceId, Long namespa
         Node node = new Node();
         node.setId(datasourceId);
         node.setType(NodeTypeEnum.DATA_SOURCE.name());
-        TreeNodeStorage.INSTANCE.updatePosition(dropToNode, node, 2);
+        TreeNodeStorage.INSTANCE.insertNode(dropToNode, node);
     }
 
     public void delete(Long id) {

chat2db-community-server/chat2db-community-storage/src/main/java/ai/chat2db/community/storage/small/NamespaceStorage.java

@@ -65,7 +65,7 @@ public Long save(Namespace namespace){
         Node node = new Node();
         node.setId(namespace.getId());
         node.setType(NodeTypeEnum.NAMESPACE.name());
-        TreeNodeStorage.INSTANCE.updatePosition(dropToNode, node, 2);
+        TreeNodeStorage.INSTANCE.insertNode(dropToNode, node);
         return id;
     }
     public void delete(Long id) {

chat2db-community-server/chat2db-community-storage/src/main/java/ai/chat2db/community/storage/small/TreeNodeStorage.java

@@ -6,25 +6,21 @@
 import ai.chat2db.community.tools.wrapper.result.ActionResult;
 import com.alibaba.fastjson2.JSON;
 import com.alibaba.fastjson2.filter.PropertyFilter;
-import com.google.common.collect.Lists;
-import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.collections4.CollectionUtils;
-import org.apache.commons.collections4.MapUtils;
 
 import java.io.File;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.TreeMap;
 
-@Slf4j
 public class TreeNodeStorage extends SmallDataStorage<TreeNode> {
     public static final TreeNodeStorage INSTANCE = new TreeNodeStorage();
 
     protected TreeNodeStorage() {
         super("tree", TreeNode.class);
-        if (MapUtils.isEmpty(dataMap)) {
-
-        }
     }
 
     TreeNodeStorage(File storageFile) {
@@ -33,130 +29,158 @@ protected TreeNodeStorage() {
 
     public synchronized List<Node> getNodes() {
         List<TreeNode> treeNodes = getDataList();
-        if (treeNodes == null) {
-            return null;
-        }
-        if (CollectionUtils.isEmpty(treeNodes)) {
-            return Lists.newArrayList();
-        }
-        return treeNodes.get(0).getChildren();
+        return treeNodes.isEmpty() ? new ArrayList<>() : treeNodes.get(0).getChildren();
     }
 
     public synchronized void createTree(List<Node> nodes) {
-        if (nodes == null) {
-            return;
+        if (nodes != null) {
+            persistTree(copyNodes(nodes));
         }
-        PropertyFilter filter = (object, name, value) -> !"data".equals(name);
-        String json = JSON.toJSONString(nodes, filter);
-        List<Node> newNodes = JSON.parseArray(json, Node.class);
+    }
 
+    private void persistTree(List<Node> nodes) {
         List<TreeNode> treeNodes = getDataList();
-        if (CollectionUtils.isEmpty(treeNodes)) {
-            TreeNode treeNode = new TreeNode();
-            treeNode.setChildren(newNodes);
-            save(treeNode);
-        } else {
-            treeNodes.get(0).setChildren(newNodes);
-            update(treeNodes.get(0));
+        TreeNode replacement = new TreeNode();
+        replacement.setId(treeNodes.isEmpty() ? generateId() : treeNodes.get(0).getId());
+        replacement.setChildren(nodes);
+        Map<Long, TreeNode> persistedData = new TreeMap<>(dataMap);
+        persistedData.put(replacement.getId(), replacement);
+        saveDataList(new ArrayList<>(persistedData.values()));
+        dataMap.put(replacement.getId(), replacement);
+    }
+
+    synchronized void insertNode(Node parentNode, Node newNode) {
+        if (newNode == null) {
+            return;
         }
+        List<Node> nodes = getNodes();
+        List<Node> updatedNodes = nodes == null ? new ArrayList<>() : copyNodes(nodes);
+        if (findNode(updatedNodes, newNode) != null) {
+            return;
+        }
+        if (parentNode == null) {
+            updatedNodes.add(newNode);
+        } else if (!addNode(updatedNodes, parentNode, newNode, 2)) {
+            return;
+        }
+        createTree(updatedNodes);
     }
 
 
     public synchronized ActionResult updatePosition(Node dropToNode, Node dragNode, Integer dropPosition) {
         if (dragNode == null) {
             return ActionResult.isSuccess();
         }
-        try {
-            Thread.sleep(10);
-        } catch (InterruptedException e) {
-            Thread.currentThread().interrupt();
-        }
         List<Node> nodes = getNodes();
         if (nodes == null) {
             return ActionResult.isSuccess();
         }
-        if (dropToNode == null) {
-            removeNode(nodes, dragNode, false);
-            nodes.add(dragNode);
-            createTree(nodes);
+        List<Node> updatedNodes = copyNodes(nodes);
+        Node sourceNode = findNode(updatedNodes, dragNode);
+        if (sourceNode == null) {
+            return ActionResult.isSuccess();
+        }
+        Node targetNode = findNode(updatedNodes, dropToNode);
+        if (dropToNode != null && (targetNode == null || sameNode(sourceNode, targetNode)
+                || findNode(sourceNode.getChildren(), targetNode) != null)) {
             return ActionResult.isSuccess();
-        } else {
-            removeNode(nodes, dragNode, false);
-            addNode(nodes, dropToNode, dragNode, dropPosition);
-            createTree(nodes);
         }
+        removeNode(updatedNodes, sourceNode, false);
+        if (targetNode == null) {
+            updatedNodes.add(sourceNode);
+        } else if (!addNode(updatedNodes, targetNode, sourceNode, dropPosition)) {
+            return ActionResult.isSuccess();
+        }
+        persistTree(updatedNodes);
         return ActionResult.isSuccess();
     }
 
-    private synchronized void removeNode(List<Node> nodes, Node dragNode, boolean deleteChildren) {
+    private List<Node> copyNodes(List<Node> nodes) {
+        PropertyFilter filter = (object, name, value) -> !"data".equals(name);
+        String json = JSON.toJSONString(nodes, filter);
+        return JSON.parseArray(json, Node.class);
+    }
+
+    private Node findNode(List<Node> nodes, Node expected) {
+        if (CollectionUtils.isEmpty(nodes) || expected == null) {
+            return null;
+        }
+        for (Node node : nodes) {
+            if (sameNode(node, expected)) {
+                return node;
+            }
+            Node child = findNode(node.getChildren(), expected);
+            if (child != null) {
+                return child;
+            }
+        }
+        return null;
+    }
+
+    private boolean sameNode(Node left, Node right) {
+        return left != null && right != null
+                && Objects.equals(left.getId(), right.getId())
+                && Objects.equals(left.getType(), right.getType());
+    }
+
+    private boolean removeNode(List<Node> nodes, Node dragNode, boolean deleteChildren) {
         if (CollectionUtils.isEmpty(nodes)) {
-            return;
+            return false;
         }
         Iterator<Node> iterator = nodes.iterator();
-        List<Node> tempList = new ArrayList<>();
         while (iterator.hasNext()) {
             Node node = iterator.next();
-            if (node.getId().equals(dragNode.getId()) && node.getType().equals(dragNode.getType())) {
-                if (NodeTypeEnum.NAMESPACE.name().equals(node.getType())) {
-                    List<Node> c = node.getChildren();
-                    if (CollectionUtils.isNotEmpty(c)) {
-                        tempList.addAll(c);
-                        dragNode.setChildren(c);
-                    }
-                }
+            if (sameNode(node, dragNode)) {
                 iterator.remove();
-                if (CollectionUtils.isNotEmpty(tempList) && deleteChildren) {
-                    nodes.addAll(tempList);
+                if (deleteChildren && NodeTypeEnum.NAMESPACE.name().equals(node.getType())
+                        && CollectionUtils.isNotEmpty(node.getChildren())) {
+                    nodes.addAll(node.getChildren());
                 }
-                return;
+                return true;
+            }
+            if (removeNode(node.getChildren(), dragNode, deleteChildren)) {
+                return true;
             }
-            removeNode(node.getChildren(), dragNode, deleteChildren);
         }
+        return false;
     }
 
     public synchronized ActionResult deleteNode(Node dragNode) {
         List<Node> nodes = getNodes();
-        removeNode(nodes, dragNode, true);
-        createTree(nodes);
+        if (nodes == null || dragNode == null) {
+            return ActionResult.isSuccess();
+        }
+        List<Node> updatedNodes = copyNodes(nodes);
+        if (!removeNode(updatedNodes, dragNode, true)) {
+            return ActionResult.isSuccess();
+        }
+        persistTree(updatedNodes);
         return ActionResult.isSuccess();
     }
 
-    private synchronized void addNode(List<Node> nodes, Node dropToNode, Node dragNode, Integer dropToGap) {
+    private boolean addNode(List<Node> nodes, Node dropToNode, Node dragNode, Integer dropToGap) {
         if (CollectionUtils.isEmpty(nodes)) {
-            return;
-        }
-        if (dropToNode.getId().equals(dragNode.getId()) && dropToNode.getType().equals(dragNode.getType())) {
-            return;
+            return false;
         }
-        int index = 0;
-        for (Node node : nodes) {
-            index++;
-            if (node.getId().equals(dropToNode.getId()) && node.getType().equals(dropToNode.getType())) {
-                if (dropToGap == 0) {
-                    List<Node> children = node.getChildren();
-                    if (children == null) {
-                        children = new ArrayList<>();
-                    }
-                    children.add(0, dragNode);
-                    node.setChildren(children);
-                    return;
-                }else if (dropToGap == 2) {
+        for (int index = 0; index < nodes.size(); index++) {
+            Node node = nodes.get(index);
+            if (sameNode(node, dropToNode)) {
+                if (dropToGap == 0 || dropToGap == 2) {
                     List<Node> children = node.getChildren();
                     if (children == null) {
                         children = new ArrayList<>();
+                        node.setChildren(children);
                     }
-                    children.add(dragNode);
-                    node.setChildren(children);
-                    return;
-                }else if (dropToGap == 1) {
-                    nodes.add(index, dragNode);
-                    return;
+                    children.add(dropToGap == 0 ? 0 : children.size(), dragNode);
                 } else {
-                    nodes.add(index - 1, dragNode);
-                    return;
+                    nodes.add(dropToGap == 1 ? index + 1 : index, dragNode);
                 }
+                return true;
+            }
+            if (addNode(node.getChildren(), dropToNode, dragNode, dropToGap)) {
+                return true;
             }
-            addNode(node.getChildren(), dropToNode, dragNode, dropToGap);
         }
+        return false;
     }
 }

chat2db-community-server/chat2db-community-storage/src/test/java/ai/chat2db/community/storage/small/TreeNodeCreationTest.java

@@ -0,0 +1,77 @@
+package ai.chat2db.community.storage.small;
+
+import ai.chat2db.community.domain.api.enums.NodeTypeEnum;
+import ai.chat2db.community.domain.api.model.storage.WorkspaceDataSource;
+import ai.chat2db.community.domain.api.model.workspace.Namespace;
+import ai.chat2db.community.domain.api.model.workspace.Node;
+import ai.chat2db.community.storage.LocalWorkspaceStorage;
+import ai.chat2db.community.storage.TestHome;
+import ai.chat2db.community.storage.converter.StorageConverterImpl;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import java.util.List;
+import static org.junit.jupiter.api.Assertions.*;
+
+class TreeNodeCreationTest {
+    private final LocalWorkspaceStorage storage = new LocalWorkspaceStorage(new StorageConverterImpl());
+
+    @BeforeAll
+    static void isolatedHome() { TestHome.init(); }
+
+    @BeforeEach
+    @AfterEach
+    void clearFixtures() {
+        TreeNodeStorage.INSTANCE.dataMap.clear();
+        TreeNodeStorage.INSTANCE.saveDataList();
+        NamespaceStorage.INSTANCE.dataMap.clear();
+        NamespaceStorage.INSTANCE.saveDataList();
+        DataSourceStorage.INSTANCE.dataMap.clear();
+        DataSourceStorage.INSTANCE.saveDataList();
+    }
+
+    @Test
+    void newRootNamespaceIsVisibleInAlreadyPopulatedTree() {
+        long parentId = rootGroup();
+        Namespace sibling = new Namespace();
+        sibling.setName("review sibling");
+        long siblingId = storage.createNamespace(sibling);
+        assertEquals(List.of(parentId, siblingId), storage.getTree().stream().map(Node::getId).toList());
+    }
+
+    @Test
+    void newChildNamespaceKeepsRequestedParent() {
+        long parentId = rootGroup();
+        Namespace child = new Namespace();
+        child.setName("review child");
+        child.setParentId(parentId);
+        long childId = storage.createNamespace(child);
+        Node parent = storage.getTree().get(0);
+        assertEquals(parentId, parent.getId());
+        assertNotNull(parent.getChildren(), "new child must be inserted into the persisted tree");
+        assertEquals(List.of(childId), parent.getChildren().stream().map(Node::getId).toList());
+    }
+
+    @Test
+    void newDatasourceIsVisibleInRequestedGroup() {
+        long parentId = rootGroup();
+        WorkspaceDataSource datasource = new WorkspaceDataSource();
+        datasource.setAlias("review datasource");
+        datasource.setSpaceId(parentId);
+        long id = storage.createDataSource(datasource);
+        assertNotNull(storage.queryDataSourceById(id, false), "datasource record was persisted");
+        Node parent = storage.getTree().get(0);
+        assertNotNull(parent.getChildren(), "new datasource must be inserted into the persisted tree");
+        assertEquals(List.of(id), parent.getChildren().stream().map(Node::getId).toList());
+        assertEquals(NodeTypeEnum.DATA_SOURCE.name(), parent.getChildren().get(0).getType());
+    }
+
+    private long rootGroup() {
+        Namespace parent = new Namespace();
+        parent.setName("review parent");
+        long id = storage.createNamespace(parent);
+        assertEquals(List.of(id), storage.getTree().stream().map(Node::getId).toList());
+        return id;
+    }
+}