Browse Source

fix: cleanup data after disconnect (#871)

Signed-off-by: ryjiang <jiangruiyi@gmail.com>
ryjiang 2 months ago
parent
commit
2351eb152d

+ 46 - 15
client/src/context/hooks/useCollectionsManagement.ts

@@ -1,20 +1,18 @@
-import { useCallback, useRef, useState, useEffect } from 'react';
+import { useCallback, useRef, useState, useEffect, useContext } from 'react';
 import { CollectionService, MilvusService } from '@/http';
 import { WS_EVENTS, WS_EVENTS_TYPE } from '@server/utils/Const';
 import { checkIndexing, checkLoading } from '@server/utils/Shared';
+import { authContext } from '@/context';
 import type { CollectionObject, CollectionFullObject } from '@server/types';
 
 export function useCollectionsManagement(database: string) {
   const [collections, setCollections] = useState<CollectionObject[]>([]);
+  const { isAuth } = useContext(authContext);
+
   const [loading, setLoading] = useState(true);
   const requestIdRef = useRef(0);
   const databaseRef = useRef(database);
 
-  // Update the ref when database changes
-  useEffect(() => {
-    databaseRef.current = database;
-  }, [database]);
-
   const detectLoadingIndexing = useCallback(
     (collections: CollectionObject[]) => {
       const LoadingOrBuildingCollections = collections.filter(v => {
@@ -102,6 +100,33 @@ export function useCollectionsManagement(database: string) {
     >
   >(new Map());
 
+  function clearRefreshCollectionsDebounceMapRef() {
+    refreshCollectionsDebounceMapRef.current.forEach(ref => {
+      if (ref.timer) {
+        clearTimeout(ref.timer);
+      }
+      ref.names = [];
+      ref.timer = null;
+    });
+    refreshCollectionsDebounceMapRef.current.clear();
+  }
+
+  useEffect(() => {
+    if (!isAuth) {
+      // clear collections when not authenticated
+      setCollections([]);
+      clearRefreshCollectionsDebounceMapRef();
+    }
+  }, [isAuth]);
+
+  // Update the ref when database changes
+  useEffect(() => {
+    databaseRef.current = database;
+    setCollections([]);
+    // clear refreshCollectionsDebounceMapRef
+    clearRefreshCollectionsDebounceMapRef();
+  }, [database]);
+
   const fetchCollections = async () => {
     const currentRequestId = ++requestIdRef.current;
     try {
@@ -122,17 +147,23 @@ export function useCollectionsManagement(database: string) {
     }
   };
 
-  const fetchCollection = async (name: string, drop?: boolean) => {
-    if (drop) {
-      updateCollections({ collections: [], deletedNames: [name] });
-    } else {
-      const res = await CollectionService.getCollection(name);
-      updateCollections({ collections: [res] });
-    }
-  };
+  const fetchCollection = useCallback(
+    async (name: string, drop?: boolean) => {
+      if (!isAuth) return;
+      if (drop) {
+        updateCollections({ collections: [], deletedNames: [name] });
+      } else {
+        const res = await CollectionService.getCollection(name);
+        updateCollections({ collections: [res] });
+      }
+    },
+    [isAuth, updateCollections]
+  );
 
   const batchRefreshCollections = useCallback(
     async (collectionNames: string[], key: string = 'default') => {
+      // if not logged in, do not refresh collections
+      if (!isAuth) return;
       let ref = refreshCollectionsDebounceMapRef.current.get(key);
       if (!ref) {
         ref = { timer: null, names: [], pending: new Set() };
@@ -181,7 +212,7 @@ export function useCollectionsManagement(database: string) {
         ref!.timer = null;
       }, 200);
     },
-    [collections, updateCollections] // Removed database dependency
+    [collections, updateCollections, isAuth] // Removed database dependency
   );
 
   return {

+ 42 - 37
server/src/collections/collections.service.ts

@@ -473,50 +473,55 @@ export class CollectionsService {
     collections: string[] = [],
     database?: string
   ): Promise<CollectionObject[]> {
-    // get all collections(name, timestamp, id)
-    const allCollections = await this.showCollections(clientId, {
-      db_name: database,
-    });
-    // get all loaded collection
-    const loadedCollections = await this.showCollections(clientId, {
-      type: ShowCollectionsType.Loaded,
-      db_name: database,
-    });
+    try {
+      // get all collections(name, timestamp, id)
+      const allCollections = await this.showCollections(clientId, {
+        db_name: database,
+      });
+      // get all loaded collection
+      const loadedCollections = await this.showCollections(clientId, {
+        type: ShowCollectionsType.Loaded,
+        db_name: database,
+      });
 
-    // data container
-    const data: CollectionObject[] = [];
+      // data container
+      const data: CollectionObject[] = [];
 
-    // get target collections details
-    const targetCollections = allCollections.data.filter(
-      d => collections.indexOf(d.name) !== -1
-    );
+      // get target collections details
+      const targetCollections = allCollections.data.filter(
+        d => collections.indexOf(d.name) !== -1
+      );
 
-    const targets =
-      targetCollections.length > 0 ? targetCollections : allCollections.data;
+      const targets =
+        targetCollections.length > 0 ? targetCollections : allCollections.data;
 
-    // sort targets by name
-    targets.sort((a, b) => a.name.localeCompare(b.name));
+      // sort targets by name
+      targets.sort((a, b) => a.name.localeCompare(b.name));
 
-    // get all collection details
-    for (let i = 0; i < targets.length; i++) {
-      const collection = targets[i];
-      const loadedCollection = loadedCollections.data.find(
-        v => v.name === collection.name
-      );
+      // get all collection details
+      for (let i = 0; i < targets.length; i++) {
+        const collection = targets[i];
+        const loadedCollection = loadedCollections.data.find(
+          v => v.name === collection.name
+        );
 
-      data.push(
-        await this.getCollection(
-          clientId,
-          collection,
-          loadedCollection,
-          collections.length === 0, // if no collection is specified, load all collections without detail
-          database
-        )
-      );
-    }
+        data.push(
+          await this.getCollection(
+            clientId,
+            collection,
+            loadedCollection,
+            collections.length === 0, // if no collection is specified, load all collections without detail
+            database
+          )
+        );
+      }
 
-    // return data
-    return data;
+      // return data
+      return data;
+    } catch (error) {
+      // if error occurs, return empty array, for example, when the client is disconnected
+      return [];
+    }
   }
 
   // update collections details