feat: add isComplete to streamed chart query results

Author: pulpdrew

packages/app/src/components/PatternTable.tsx

@@ -29,11 +29,11 @@ export default function PatternTable({
 
   const [selectedPattern, setSelectedPattern] = useState<Pattern | null>(null);
 
-  const { totalCount, isLoading: isTotalCountLoading } = useSearchTotalCount(
-    totalCountConfig,
-    totalCountQueryKeyPrefix,
-    true, // disable chunked queries for total count to avoid estimated counts flickering
-  );
+  const {
+    totalCount,
+    isLoading: isTotalCountLoading,
+    isTotalCountComplete,
+  } = useSearchTotalCount(totalCountConfig, totalCountQueryKeyPrefix);
 
   const {
     data: groupedResults,
@@ -47,7 +47,8 @@ export default function PatternTable({
     totalCount,
   });
 
-  const isLoading = isTotalCountLoading || isGroupedPatternsLoading;
+  const isLoading =
+    isTotalCountLoading || !isTotalCountComplete || isGroupedPatternsLoading;
 
   const sortedGroupedResults = useMemo(() => {
     return Object.values(groupedResults).sort(

packages/app/src/components/SearchTotalCountChart.tsx

@@ -9,7 +9,6 @@ import { useQueriedChartConfig } from '@/hooks/useChartConfig';
 export function useSearchTotalCount(
   config: ChartConfigWithDateRange,
   queryKeyPrefix: string,
-  disableQueryChunking: boolean = false,
 ) {
   // copied from DBTimeChart
   const { granularity } = useTimeChartSettings(config);
@@ -23,13 +22,14 @@ export function useSearchTotalCount(
     isLoading,
     isError,
   } = useQueriedChartConfig(queriedConfig, {
-    queryKey: [queryKeyPrefix, queriedConfig, disableQueryChunking],
+    queryKey: [queryKeyPrefix, queriedConfig],
     staleTime: 1000 * 60 * 5,
     refetchOnWindowFocus: false,
     placeholderData: keepPreviousData, // no need to flash loading state when in live tail
-    disableQueryChunking,
   });
 
+  const isTotalCountComplete = !!totalCountData?.isComplete;
+
   const totalCount = useMemo(() => {
     return totalCountData?.data?.reduce(
       (p: number, v: any) => p + Number.parseInt(v['count()']),
@@ -41,6 +41,7 @@ export function useSearchTotalCount(
     totalCount,
     isLoading,
     isError,
+    isTotalCountComplete,
   };
 }
 

packages/app/src/hooks/__tests__/useChartConfig.test.tsx

@@ -357,6 +357,7 @@ describe('useChartConfig', () => {
         data: mockResponse.data,
         meta: mockResponse.meta,
         rows: mockResponse.rows,
+        isComplete: true,
       });
       expect(result.current.isLoading).toBe(false);
       expect(result.current.isPending).toBe(false);
@@ -402,6 +403,7 @@ describe('useChartConfig', () => {
         data: mockResponse.data,
         meta: mockResponse.meta,
         rows: mockResponse.rows,
+        isComplete: true,
       });
       expect(result.current.isLoading).toBe(false);
       expect(result.current.isPending).toBe(false);
@@ -452,6 +454,7 @@ describe('useChartConfig', () => {
         data: mockResponse.data,
         meta: mockResponse.meta,
         rows: mockResponse.rows,
+        isComplete: true,
       });
     });
 
@@ -527,6 +530,7 @@ describe('useChartConfig', () => {
         data: mockResponse.data,
         meta: mockResponse.meta,
         rows: mockResponse.rows,
+        isComplete: true,
       });
     });
 
@@ -577,6 +581,7 @@ describe('useChartConfig', () => {
         data: mockResponse.data,
         meta: mockResponse.meta,
         rows: mockResponse.rows,
+        isComplete: true,
       });
     });
 
@@ -667,6 +672,7 @@ describe('useChartConfig', () => {
         ],
         meta: mockResponse1.meta,
         rows: 5,
+        isComplete: true,
       });
       expect(result.current.isLoading).toBe(false);
       expect(result.current.isPending).toBe(false);
@@ -726,6 +732,7 @@ describe('useChartConfig', () => {
         data: [...mockResponse2.data, ...mockResponse1.data],
         meta: mockResponse1.meta,
         rows: 4,
+        isComplete: false,
       });
       expect(result.current.isFetching).toBe(true);
       expect(result.current.isLoading).toBe(false); // isLoading is false because we have partial data
@@ -750,6 +757,7 @@ describe('useChartConfig', () => {
         ],
         meta: mockResponse1.meta,
         rows: 5,
+        isComplete: true,
       });
     });
 
@@ -798,6 +806,7 @@ describe('useChartConfig', () => {
         data: mockResponse1.data,
         meta: mockResponse1.meta,
         rows: 1,
+        isComplete: false,
       });
     });
 

packages/app/src/hooks/useChartConfig.tsx

@@ -48,7 +48,9 @@ type TimeWindow = {
   dateRangeEndInclusive?: boolean;
 };
 
-type TQueryFnData = Pick<ResponseJSON<any>, 'data' | 'meta' | 'rows'> & {};
+type TQueryFnData = Pick<ResponseJSON<any>, 'data' | 'meta' | 'rows'> & {
+  isComplete: boolean;
+};
 
 const shouldUseChunking = (
   config: ChartConfigWithOptDateRange,
@@ -125,7 +127,6 @@ async function* fetchDataInChunks(
       ? getGranularityAlignedTimeWindows(config)
       : ([undefined] as const);
 
-  let query = null;
   if (IS_MTVIEWS_ENABLED) {
     const { dataTableDDL, mtViewDDL, renderMTViewConfig } =
       await buildMTViewSelectQuery(config);
@@ -134,10 +135,12 @@ async function* fetchDataInChunks(
     console.log('dataTableDDL:', dataTableDDL);
     // eslint-disable-next-line no-console
     console.log('mtViewDDL:', mtViewDDL);
-    query = await renderMTViewConfig();
+    await renderMTViewConfig();
   }
 
-  for (const window of windows) {
+  for (let i = 0; i < windows.length; i++) {
+    const window = windows[i];
+
     const windowedConfig = {
       ...config,
       ...(window ?? {}),
@@ -151,7 +154,7 @@ async function* fetchDataInChunks(
       },
     });
 
-    yield result;
+    yield { chunk: result, isComplete: i === windows.length - 1 };
   }
 }
 
@@ -164,13 +167,17 @@ async function* fetchDataInChunks(
  *
  * For chunked queries, note the following:
  * - `config.limit`, if provided, is applied to each chunk, so the total number
- *   of rows returned may be up to `limit * number_of_chunks`.
+ *    of rows returned may be up to `limit * number_of_chunks`.
  * - The returned data will be ordered within each chunk, and chunks will
- *   be ordered oldest-first, by the `timestampValueExpression`.
+ *    be ordered oldest-first, by the `timestampValueExpression`.
+ * - `isPending` is true until the first chunk is fetched. Once the first chunk
+ *    is available, `isPending` will be false and `isSuccess` will be true.
+ *    `isFetching` will be true until all chunks have been fetched.
+ * - `data.isComplete` indicates whether all chunks have been fetched.
  */
 export function useQueriedChartConfig(
   config: ChartConfigWithOptDateRange,
-  options?: Partial<UseQueryOptions<ResponseJSON<any>>> &
+  options?: Partial<UseQueryOptions<TQueryFnData>> &
     AdditionalUseQueriedChartConfigOptions,
 ) {
   const clickhouseClient = useClickhouseClient();
@@ -194,12 +201,18 @@ export function useQueriedChartConfig(
        * in flickering or render loops.
        */
       refetchMode: 'replace',
-      initialValue: { data: [], meta: [], rows: 0 } as TQueryFnData,
-      reducer: (acc, chunk) => {
+      initialValue: {
+        data: [],
+        meta: [],
+        rows: 0,
+        isComplete: false,
+      } as TQueryFnData,
+      reducer: (acc, { chunk, isComplete }) => {
         return {
           data: [...(chunk.data || []), ...(acc?.data || [])],
           meta: chunk.meta,
           rows: (acc?.rows || 0) + (chunk.rows || 0),
+          isComplete,
         };
       },
     }),

packages/app/src/hooks/usePatterns.tsx

@@ -141,14 +141,15 @@ function usePatterns({
     limit: { limit: samples },
   });
 
-  const { data: sampleRows } = useQueriedChartConfig(
-    configWithPrimaryAndPartitionKey ?? config, // `config` satisfying type, never used due to `enabled` check
-    {
-      enabled: configWithPrimaryAndPartitionKey != null && enabled,
-      // Disable chunking to ensure we get the desired sample size
-      disableQueryChunking: true,
-    },
-  );
+  const { data: sampleRows, isLoading: isSampleLoading } =
+    useQueriedChartConfig(
+      configWithPrimaryAndPartitionKey ?? config, // `config` satisfying type, never used due to `enabled` check
+      {
+        enabled: configWithPrimaryAndPartitionKey != null && enabled,
+        // Disable chunking to ensure we get the desired sample size
+        disableQueryChunking: true,
+      },
+    );
 
   const { data: pyodide, isLoading: isLoadingPyodide } = usePyodide({
     enabled,
@@ -195,7 +196,7 @@ function usePatterns({
 
   return {
     ...query,
-    isLoading: query.isLoading || isLoadingPyodide,
+    isLoading: query.isLoading || isSampleLoading || isLoadingPyodide,
     patternQueryConfig: configWithPrimaryAndPartitionKey,
   };
 }