feat: keep query parameters when switching pages (#746)

Signed-off-by: ryjiang <jiangruiyi@gmail.com>

client/src/hooks/Query.ts

@@ -1,36 +1,35 @@
 import { useState, useRef, useEffect } from 'react';
 import { DataTypeStringEnum, MIN_INT64 } from '@/consts';
 import { CollectionService } from '@/http';
-import type { CollectionFullObject, FieldObject } from '@server/types';
+import type { CollectionFullObject } from '@server/types';
 
 // TODO: refactor this, a little bit messy
 export const useQuery = (params: {
   collection: CollectionFullObject;
-  fields: FieldObject[];
+  outputFields: string[];
   consistencyLevel: string;
   onQueryStart: Function;
   onQueryEnd?: Function;
   onQueryFinally?: Function;
+  initialExpr: string;
 }) => {
   // params
   const {
     collection,
-    fields,
+    outputFields,
     onQueryStart,
     onQueryEnd,
     onQueryFinally,
     consistencyLevel,
+    initialExpr,
   } = params;
 
   // states
   const [currentPage, setCurrentPage] = useState<number>(0);
   const [pageSize, setPageSize] = useState<number>(0);
   const [total, setTotal] = useState<number>(collection.rowCount);
-  const [expr, setExpr] = useState<string>('');
+  const [expr, setExpr] = useState<string>(initialExpr);
   const [queryResult, setQueryResult] = useState<any>({ data: [], latency: 0 });
-  const [outputFields, setOutputFields] = useState<string[]>(
-    fields.map(i => i.name) || []
-  );
 
   // build local cache for pk ids
   const pageCache = useRef(new Map());
@@ -159,13 +158,14 @@ export const useQuery = (params: {
     // get count;
     count();
 
-    // update output fields, then query again
-    const newOutputFields = fields.map(i => i.name) || [];
-    setOutputFields(newOutputFields);
-
     // do the query
-    query(currentPage, consistencyLevel, newOutputFields);
-  }, [expr, pageSize, collection.collection_name]);
+    query(currentPage, consistencyLevel, outputFields);
+  }, [
+    expr,
+    pageSize,
+    collection.collection_name,
+    JSON.stringify(outputFields),
+  ]);
 
   return {
     // total query count
@@ -192,9 +192,5 @@ export const useQuery = (params: {
     count,
     // get expression
     getPageExpr,
-    // output fields
-    outputFields,
-    // set output fields
-    setOutputFields,
   };
 };

client/src/pages/databases/Databases.tsx

@@ -6,7 +6,6 @@ import { useNavigationHook } from '@/hooks';
 import { ALL_ROUTER_TYPES } from '@/router/consts';
 import RouteTabList from '@/components/customTabList/RouteTabList';
 import DatabaseTree from '@/pages/databases/tree';
-import { ITab } from '@/components/customTabList/Types';
 import Partitions from './collections/partitions/Partitions';
 import Schema from './collections/schema/Schema';
 import Data from './collections/data/CollectionData';
@@ -17,9 +16,10 @@ import { dataContext, authContext } from '@/context';
 import Collections from './collections/Collections';
 import StatusIcon, { LoadingType } from '@/components/status/StatusIcon';
 import { ConsistencyLevelEnum, DYNAMIC_FIELD } from '@/consts';
-import { SearchParams } from './types';
-import { CollectionObject, CollectionFullObject } from '@server/types';
 import { makeStyles } from '@mui/styles';
+import type { SearchParams, QueryState } from './types';
+import type { CollectionObject, CollectionFullObject } from '@server/types';
+import type { ITab } from '@/components/customTabList/Types';
 
 const DEFAULT_TREE_WIDTH = 230;
 
@@ -86,9 +86,8 @@ const Databases = () => {
     useContext(dataContext);
 
   // UI state
-  const [searchParams, setSearchParams] = useState<SearchParams[]>(
-    [] as SearchParams[]
-  );
+  const [searchParams, setSearchParams] = useState<SearchParams[]>([]);
+  const [queryState, setQueryState] = useState<QueryState[]>([]);
 
   // tree ref
   const [isDragging, setIsDragging] = useState(false);
@@ -148,7 +147,7 @@ const Databases = () => {
     setUIPref({ tree: { width: DEFAULT_TREE_WIDTH } });
   };
 
-  // init search params
+  // init search params and query state
   useEffect(() => {
     collections.forEach(c => {
       // find search params for the collection
@@ -218,6 +217,40 @@ const Databases = () => {
           });
         });
       }
+
+      // find query state for the collection
+      const query = queryState.find(
+        q => q.collection.collection_name === c.collection_name
+      );
+
+      // if query state not found, and the schema is ready, create new query state
+      if (!query && c.schema) {
+        setQueryState(prevState => {
+          const fields = [...c.schema.fields, ...c.schema.dynamicFields];
+          return [
+            ...prevState,
+            {
+              collection: c,
+              expr: '',
+              fields: fields,
+              outputFields: fields.map(f => f.name),
+              consistencyLevel: ConsistencyLevelEnum.Bounded,
+            },
+          ];
+        });
+      } else {
+        // update collection
+        setQueryState(prevState => {
+          return prevState.map(q => {
+            if (q.collection.collection_name === c.collection_name) {
+              // update collection
+              const collection = c;
+              return { ...q, collection };
+            }
+            return q;
+          });
+        });
+      }
     });
 
     // delete search params for the collection that is not in the collections
@@ -228,6 +261,15 @@ const Databases = () => {
         )
       );
     });
+
+    // delete query state for the collection that is not in the collections
+    setQueryState(prevState => {
+      return prevState.filter(q =>
+        collections.find(
+          c => c.collection_name === q.collection.collection_name
+        )
+      );
+    });
   }, [collections]);
 
   // get current collection from url
@@ -261,6 +303,17 @@ const Databases = () => {
     });
   };
 
+  const setCollectionQueryState = (state: QueryState) => {
+    setQueryState(prevState => {
+      return prevState.map(q => {
+        if (q.collection.collection_name === state.collection.collection_name) {
+          return { ...state };
+        }
+        return q;
+      });
+    });
+  };
+
   // render
   return (
     <section
@@ -308,6 +361,12 @@ const Databases = () => {
             )!
           }
           setSearchParams={setCollectionSearchParams}
+          queryState={
+            queryState.find(
+              q => q.collection.collection_name === collectionName
+            )!
+          }
+          setQueryState={setCollectionQueryState}
           collections={collections}
         />
       )}
@@ -361,6 +420,8 @@ const CollectionTabs = (props: {
   collections: CollectionObject[]; // collections
   searchParams: SearchParams; // search params
   setSearchParams: (params: SearchParams) => void; // set search params
+  queryState: QueryState; // query state
+  setQueryState: (state: QueryState) => void; // set query state
 }) => {
   // props
   const {
@@ -370,6 +431,8 @@ const CollectionTabs = (props: {
     collections,
     searchParams,
     setSearchParams,
+    queryState,
+    setQueryState,
   } = props;
 
   // context
@@ -402,9 +465,7 @@ const CollectionTabs = (props: {
     },
     {
       label: collectionTrans('dataTab'),
-      component: (
-        <Data collections={collections} collectionName={collectionName} />
-      ),
+      component: <Data queryState={queryState} setQueryState={setQueryState} />,
       path: `data`,
     },
     {

client/src/pages/databases/collections/data/CollectionData.tsx

@@ -30,19 +30,17 @@ import CustomMultiSelector from '@/components/customSelector/CustomMultiSelector
 import CollectionColHeader from '../CollectionColHeader';
 import DataView from '@/components/DataView/DataView';
 import DataListView from '@/components/DataListView/DataListView';
-import type { CollectionObject, CollectionFullObject } from '@server/types';
+import type { QueryState } from '../../types';
 
 export interface CollectionDataProps {
-  collectionName: string;
-  collections: CollectionObject[];
+  queryState: QueryState;
+  setQueryState: (state: QueryState) => void;
 }
 
 const CollectionData = (props: CollectionDataProps) => {
   // props
-  const { collections } = props;
-  const collection = collections.find(
-    i => i.collection_name === props.collectionName
-  ) as CollectionFullObject;
+  const { queryState, setQueryState } = props;
+  const collection = queryState && queryState.collection;
 
   // collection is not found or collection full object is not ready
   if (!collection || !collection.consistency_level) {
@@ -52,16 +50,6 @@ const CollectionData = (props: CollectionDataProps) => {
   // UI state
   const [tableLoading, setTableLoading] = useState<boolean>();
   const [selectedData, setSelectedData] = useState<any[]>([]);
-  const [expression, setExpression] = useState<string>('');
-  const [consistencyLevel, setConsistencyLevel] = useState<string>(
-    collection.consistency_level
-  );
-
-  // collection fields, combine static and dynamic fields
-  const fields = [
-    ...collection.schema.fields,
-    ...collection.schema.dynamicFields,
-  ];
 
   // UI functions
   const { setDialog, handleCloseDialog, openSnackBar, setDrawer } =
@@ -90,7 +78,7 @@ const CollectionData = (props: CollectionDataProps) => {
     const currentFilter: any = filterRef.current;
     currentFilter?.getReset();
     // update UI expression
-    setExpression('');
+    setQueryState({ ...queryState, expr: '' });
     // reset query
     reset();
     // ensure not loading
@@ -98,13 +86,13 @@ const CollectionData = (props: CollectionDataProps) => {
   };
   const handleFilterSubmit = async (expression: string) => {
     // update UI expression
-    setExpression(expression);
+    setQueryState({ ...queryState, expr: expression });
     // update expression
     setExpr(expression);
   };
   const handlePageChange = async (e: any, page: number) => {
     // do the query
-    await query(page, consistencyLevel);
+    await query(page, queryState.consistencyLevel);
     // update page number
     setCurrentPage(page);
   };
@@ -140,7 +128,6 @@ const CollectionData = (props: CollectionDataProps) => {
     currentPage,
     total,
     pageSize,
-    expr,
     queryResult,
     setPageSize,
     setCurrentPage,
@@ -148,12 +135,11 @@ const CollectionData = (props: CollectionDataProps) => {
     query,
     reset,
     count,
-    outputFields,
-    setOutputFields,
+    expr,
   } = useQuery({
     collection,
-    consistencyLevel,
-    fields: fields.filter(f => !f.is_function_output),
+    consistencyLevel: queryState.consistencyLevel,
+    outputFields: queryState.outputFields,
     onQueryStart: (expr: string = '') => {
       setTableLoading(true);
       if (expr === '') {
@@ -164,6 +150,7 @@ const CollectionData = (props: CollectionDataProps) => {
     onQueryFinally: () => {
       setTableLoading(false);
     },
+    initialExpr: queryState.expr,
   });
 
   const onInsert = async (collectionName: string) => {
@@ -380,15 +367,15 @@ const CollectionData = (props: CollectionDataProps) => {
               <CustomInput
                 type="text"
                 textConfig={{
-                  label: expression
+                  label: queryState.expr
                     ? collectionTrans('queryExpression')
                     : collectionTrans('exprPlaceHolder'),
                   key: 'advFilter',
                   className: 'textarea',
                   onChange: (value: string) => {
-                    setExpression(value);
+                    setQueryState({ ...queryState, expr: value });
                   },
-                  value: expression,
+                  value: queryState.expr,
                   disabled: !collection.loaded,
                   variant: 'filled',
                   required: false,
@@ -409,8 +396,8 @@ const CollectionData = (props: CollectionDataProps) => {
                     if (e.key === 'Enter') {
                       // reset page
                       setCurrentPage(0);
-                      if (expr !== expression) {
-                        setExpr(expression);
+                      if (expr !== queryState.expr) {
+                        setExpr(queryState.expr);
                       } else {
                         // ensure query
                         query();
@@ -424,14 +411,17 @@ const CollectionData = (props: CollectionDataProps) => {
 
               <CustomSelector
                 options={CONSISTENCY_LEVEL_OPTIONS}
-                value={consistencyLevel}
+                value={queryState.consistencyLevel}
                 label={collectionTrans('consistency')}
                 wrapperClass={classes.selector}
                 disabled={!collection.loaded}
                 variant="filled"
                 onChange={(e: { target: { value: unknown } }) => {
                   const consistency = e.target.value as string;
-                  setConsistencyLevel(consistency);
+                  setQueryState({
+                    ...queryState,
+                    consistencyLevel: consistency,
+                  });
                 }}
               />
             </div>
@@ -439,7 +429,7 @@ const CollectionData = (props: CollectionDataProps) => {
             <div className="right">
               <CustomMultiSelector
                 className={classes.outputs}
-                options={fields
+                options={queryState.fields
                   .filter(f => !f.is_function_output)
                   .map(f => {
                     return {
@@ -450,7 +440,7 @@ const CollectionData = (props: CollectionDataProps) => {
                       value: f.name,
                     };
                   })}
-                values={outputFields}
+                values={queryState.outputFields}
                 renderValue={selected => (
                   <span>{`${(selected as string[]).length} ${
                     gridTrans[
@@ -463,7 +453,7 @@ const CollectionData = (props: CollectionDataProps) => {
                 variant="filled"
                 onChange={(e: { target: { value: unknown } }) => {
                   // add value to output fields if not exist, remove if exist
-                  const newOutputFields = [...outputFields];
+                  const newOutputFields = [...queryState.outputFields];
                   const values = e.target.value as string[];
                   const newFields = values.filter(
                     v => !newOutputFields.includes(v as string)
@@ -480,11 +470,14 @@ const CollectionData = (props: CollectionDataProps) => {
                   // sort output fields by schema order
                   newOutputFields.sort(
                     (a, b) =>
-                      fields.findIndex(f => f.name === a) -
-                      fields.findIndex(f => f.name === b)
+                      queryState.fields.findIndex(f => f.name === a) -
+                      queryState.fields.findIndex(f => f.name === b)
                   );
 
-                  setOutputFields(newOutputFields);
+                  setQueryState({
+                    ...queryState,
+                    outputFields: newOutputFields,
+                  });
                 }}
               />
               <CustomButton
@@ -500,8 +493,8 @@ const CollectionData = (props: CollectionDataProps) => {
                 variant="contained"
                 onClick={() => {
                   setCurrentPage(0);
-                  if (expr !== expression) {
-                    setExpr(expression);
+                  if (expr !== queryState.expr) {
+                    setExpr(queryState.expr);
                   } else {
                     // ensure query
                     query();
@@ -515,7 +508,7 @@ const CollectionData = (props: CollectionDataProps) => {
           </div>
           <AttuGrid
             toolbarConfigs={[]}
-            colDefinitions={outputFields.map(i => {
+            colDefinitions={queryState.outputFields.map(i => {
               return {
                 id: i,
                 align: 'left',

client/src/pages/databases/types.ts

@@ -1,4 +1,8 @@
-import type { FieldObject, CollectionObject, PartitionData } from '@server/types';
+import type {
+  FieldObject,
+  CollectionObject,
+  PartitionData,
+} from '@server/types';
 
 export type SearchSingleParams = {
   anns_field: string;
@@ -58,3 +62,11 @@ export type SearchParams = {
   graphData: GraphData;
   searchLatency: number;
 };
+
+export type QueryState = {
+  collection: CollectionObject;
+  expr: string;
+  consistencyLevel: string;
+  fields: FieldObject[];
+  outputFields: string[];
+};