import { useState, useRef, useEffect } from 'react'; import { DataTypeStringEnum, DYNAMIC_FIELD, LOAD_STATE, MIN_INT64, } from '@/consts'; import { Collection } from '@/http'; export const useQuery = (params: { onQueryStart: Function; onQueryEnd?: Function; onQueryFinally?: Function; collectionName: string; }) => { // state const [collection, setCollection] = useState({ fields: [], primaryKey: { value: '', type: DataTypeStringEnum.Int64 }, loaded: false, data: null, }); const [consistencyLevel, setConsistencyLevel] = useState('Bounded'); const [currentPage, setCurrentPage] = useState(0); const [pageSize, setPageSize] = useState(0); const [total, setTotal] = useState(0); const [expr, setExpr] = useState(''); const [queryResult, setQueryResult] = useState({ data: [], latency: 0 }); const lastQuery = useRef(); // build local cache for pk ids const pageCache = useRef(new Map()); // get next/previous expression const getPageExpr = (page: number) => { const cache = pageCache.current.get( page > currentPage ? currentPage : page ); const primaryKey = collection.primaryKey; const formatPKValue = (pkId: string | number) => primaryKey.type === DataTypeStringEnum.VarChar ? `'${pkId}'` : pkId; // If cache does not exist, return expression based on primaryKey type let condition = ''; if (!cache) { const defaultValue = primaryKey.type === DataTypeStringEnum.VarChar ? "''" : `${MIN_INT64}`; condition = `${primaryKey.value} > ${defaultValue}`; } else { const { firstPKId, lastPKId } = cache; const lastPKValue = formatPKValue(lastPKId); const firstPKValue = formatPKValue(firstPKId); condition = page > currentPage ? `(${primaryKey.value} > ${lastPKValue})` : `((${primaryKey.value} <= ${lastPKValue}) && (${primaryKey.value} >= ${firstPKValue}))`; } return expr ? `${expr} && ${condition}` : condition; }; // query function const query = async ( page: number = currentPage, consistency_level = consistencyLevel ) => { const _expr = getPageExpr(page); // console.log('query expr', _expr); params.onQueryStart(_expr); try { const queryParams = { expr: _expr, output_fields: collection.fields.map((i: any) => i.name), limit: pageSize || 10, consistency_level, // travel_timestamp: timeTravelInfo.timestamp, }; // cache last query lastQuery.current = queryParams; // execute query const res = await Collection.queryData( params.collectionName, queryParams ); // get last item of the data const lastItem = res.data[res.data.length - 1]; // get first item of the data const firstItem = res.data[0]; // get last pk id const lastPKId: string | number = lastItem && lastItem[collection.primaryKey.value]; // get first pk id const firstPKId: string | number = firstItem && firstItem[collection.primaryKey.value]; // store pk id in the cache with the page number if (lastItem) { pageCache.current.set(page, { firstPKId, lastPKId, }); } // console.log('query result page', page, pageCache.current); // update query result setQueryResult(res); params.onQueryEnd?.(res); } catch (e: any) { reset(); } finally { params.onQueryFinally?.(); } }; // get collection info const prepare = async (collectionName: string) => { const collection = await Collection.getCollectionInfo(collectionName); const schemaList = collection.fields; const nameList = schemaList.map(v => ({ name: v.name, type: v.fieldType, })); // if the dynamic field is enabled, we add $meta column in the grid if (collection.enableDynamicField) { nameList.push({ name: DYNAMIC_FIELD, type: DataTypeStringEnum.JSON, }); } const primaryKey = schemaList.find(v => v.isPrimaryKey === true)!; setConsistencyLevel(collection.consistency_level); setCollection({ fields: nameList as any[], primaryKey: { value: primaryKey['name'], type: primaryKey['fieldType'] }, loaded: collection.state === LOAD_STATE.LoadStateLoaded, data: collection, }); }; const count = async (consistency_level = consistencyLevel) => { const count = 'count(*)'; const res = await Collection.queryData(params.collectionName, { expr: expr, output_fields: [count], consistency_level, // travel_timestamp: timeTravelInfo.timestamp, }); setTotal(Number(res.data[0][count])); }; // reset const reset = () => { setCurrentPage(0); setQueryResult({ data: [] }); pageCache.current.clear(); }; // Get fields at first or collection name changed. useEffect(() => { params.collectionName && prepare(params.collectionName); }, [params.collectionName]); // query if expr is changed useEffect(() => { if (!collection.primaryKey.value || !collection.loaded) { // console.info('[skip running query]: no key yet'); return; } // reset reset(); // get count; count(); // do the query query(); }, [expr]); // query if collection is changed useEffect(() => { if (!collection.primaryKey.value || !collection.loaded) { // console.info('[skip running query]: no key yet'); return; } // reset reset(); // get count; count(); // do the query query(); }, [collection]); // query if page size is changed useEffect(() => { if (!collection.primaryKey.value || !collection.loaded) { // console.info('[skip running query]: no key yet'); return; } // reset reset(); // do the query query(); }, [pageSize]); return { // collection info(primaryKey, consistency level, fields) collection, // total query count total, // page size pageSize, // update page size setPageSize, // consistency level consistencyLevel, // update consistency level setConsistencyLevel, // current page currentPage, // query current data page setCurrentPage, // expression to query expr, // expression updater setExpr, // query result queryResult, // query query, // reset reset, // count count, // get expression getPageExpr, }; };