use count api to show count for loaded collection

Signed-off-by: ryjiang <[email protected]>
zilliztech · Nov 13, 2023 · 01ba68e · 01ba68e
1 parent 23b53ea
commit 01ba68e
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 38 deletions.
diff --git a/client/src/http/Collection.ts b/client/src/http/Collection.ts
@@ -107,6 +107,13 @@ export class CollectionHttp extends BaseModel implements CollectionView {
     });
   }
 
+  static count(collectionName: string) {
+    return super.search({
+      path: `${this.COLLECTIONS_URL}/${collectionName}/count`,
+      params: {},
+    });
+  }
+
   static getQSegments(collectionName: string) {
     return super.search({
       path: `${this.COLLECTIONS_URL}/${collectionName}/qsegments`,

diff --git a/client/src/pages/overview/collectionCard/CollectionCard.tsx b/client/src/pages/overview/collectionCard/CollectionCard.tsx
@@ -6,17 +6,19 @@ import {
   Card,
   CardContent,
 } from '@material-ui/core';
-import { FC, useContext } from 'react';
+import { FC, useContext, useEffect, useState, useCallback } from 'react';
 import CustomButton from '@/components/customButton/CustomButton';
 import icons from '@/components/icons/Icons';
 import Status from '@/components/status/Status';
 import { useTranslation } from 'react-i18next';
 import CustomIconButton from '@/components/customButton/CustomIconButton';
 import { useNavigate, Link } from 'react-router-dom';
 import { LOADING_STATE } from '@/consts';
-import { rootContext } from '@/context';
+import { rootContext, dataContext } from '@/context';
 import ReleaseCollectionDialog from '../../dialogs/ReleaseCollectionDialog';
 import { CollectionCardProps } from './Types';
+import { CollectionHttp } from '@/http';
+import { CollectionData } from '@/pages/collections/Types';
 
 const useStyles = makeStyles((theme: Theme) => ({
   wrapper: {
@@ -86,16 +88,13 @@ const CollectionCard: FC<CollectionCardProps> = ({
   onRelease,
   wrapperClass = '',
 }) => {
+  const { database } = useContext(dataContext);
+  const [loading, setLoading] = useState(false);
+  const [count, setCount] = useState<string>('');
   const classes = useStyles();
   const { setDialog } = useContext(rootContext);
 
-  const {
-    _name: name,
-    _status: status,
-    _rowCount: rowCount,
-    _loadedPercentage,
-    _replicas,
-  } = data;
+  const { _name: name, _status: status, _loadedPercentage, _replicas } = data;
   const navigate = useNavigate();
   // icons
   const RightArrowIcon = icons.rightArrow;
@@ -122,6 +121,21 @@ const CollectionCard: FC<CollectionCardProps> = ({
     navigate({ pathname: '/search', search: `?collectionName=${name}` });
   };
 
+  const fetchData = useCallback(async () => {
+    try {
+      setLoading(true);
+      const data = (await CollectionHttp.count(name)) as CollectionData;
+      setCount(data._rowCount);
+    } catch (e) {
+    } finally {
+      setLoading(false);
+    }
+  }, [status]);
+
+  useEffect(() => {
+    fetchData();
+  }, [fetchData, database]);
+
   return (
     <Card
       className={`card-wrapper ${classes.wrapper} ${wrapperClass} ${
@@ -147,7 +161,11 @@ const CollectionCard: FC<CollectionCardProps> = ({
           ) : null}
           <li>
             <Typography>{collectionTrans('count')}</Typography>:
-            <Typography className={classes.rowCount}>{rowCount}</Typography>
+            {loading ? (
+              `...`
+            ) : (
+              <Typography className={classes.rowCount}>{count}</Typography>
+            )}
           </li>
         </ul>
         <Divider classes={{ root: classes.divider }} />

diff --git a/server/src/collections/collections.controller.ts b/server/src/collections/collections.controller.ts
@@ -51,6 +51,7 @@ export class CollectionController {
     );
     this.router.delete('/:name/alias/:alias', this.dropAlias.bind(this));
     this.router.get('/:name', this.describeCollection.bind(this));
+    this.router.get('/:name/count', this.count.bind(this));
 
     // load / release
     this.router.put('/:name/load', this.loadCollection.bind(this));
@@ -370,4 +371,17 @@ export class CollectionController {
       next(error);
     }
   }
+
+  async count(req: Request, res: Response, next: NextFunction) {
+    const name = req.params?.name;
+    try {
+      const result = await this.collectionsService.count({
+        collection_name: name,
+      });
+
+      res.send({ collection_name: name, rowCount: result });
+    } catch (error) {
+      next(error);
+    }
+  }
 }
diff --git a/server/src/collections/collections.service.ts b/server/src/collections/collections.service.ts
@@ -77,9 +77,15 @@ export class CollectionsService {
   }
 
   async count(data: CountReq) {
-    const res = await this.milvusService.client.count(data);
-    throwErrorFromSDK(res.status);
-    return res;
+    let count = 0;
+    try {
+      const countRes = await this.milvusService.client.count(data);
+      count = countRes.data;
+    } catch (error) {
+      const collectionStatisticsRes = await this.getCollectionStatistics(data);
+      count = collectionStatisticsRes.data.row_count;
+    }
+    return count;
   }
 
   async insert(data: InsertReq) {
@@ -172,17 +178,17 @@ export class CollectionsService {
 
         let count: number | string;
 
-        try {
-          const countRes = await this.count({
-            collection_name: name,
-          });
-          count = countRes.data;
-        } catch (error) {
-          const collectionStatisticsRes = await this.getCollectionStatistics({
-            collection_name: name,
-          });
-          count = collectionStatisticsRes.data.row_count;
-        }
+        const collectionStatisticsRes = await this.getCollectionStatistics({
+          collection_name: name,
+        });
+        count = collectionStatisticsRes.data.row_count;
+        // try {
+        //   const countRes = await this.count({
+        //     collection_name: name,
+        //   });
+        //   count = countRes.data;
+        // } catch (error) {
+        // }
 
         const indexRes = await this.getIndexInfo({
           collection_name: item.name,
@@ -241,20 +247,7 @@ export class CollectionsService {
       for (const item of res.data) {
         const { id, name } = item;
 
-        let count: number | string;
-
-        try {
-          const countRes = await this.count({
-            collection_name: name,
-          });
-          count = countRes.data;
-        } catch (error) {
-          const collectionStatisticsRes = await this.getCollectionStatistics({
-            collection_name: name,
-          });
-          count = collectionStatisticsRes.data.row_count;
-        }
-
+        const count = this.count({ collection_name: name });
         data.push({
           id,
           collection_name: name,