fix: disable query chunking for total count hook in patterns table

Author: pulpdrew

packages/app/src/components/PatternTable.tsx

@@ -32,6 +32,7 @@ export default function PatternTable({
   const { totalCount, isLoading: isTotalCountLoading } = useSearchTotalCount(
     totalCountConfig,
     totalCountQueryKeyPrefix,
+    true, // disable chunked queries for total count to avoid estimated counts flickering
   );
 
   const {

packages/app/src/components/SearchTotalCountChart.tsx

@@ -9,6 +9,7 @@ import { useQueriedChartConfig } from '@/hooks/useChartConfig';
 export function useSearchTotalCount(
   config: ChartConfigWithDateRange,
   queryKeyPrefix: string,
+  disableQueryChunking: boolean = false,
 ) {
   // copied from DBTimeChart
   const { granularity } = useTimeChartSettings(config);
@@ -22,10 +23,11 @@ export function useSearchTotalCount(
     isLoading,
     isError,
   } = useQueriedChartConfig(queriedConfig, {
-    queryKey: [queryKeyPrefix, queriedConfig],
+    queryKey: [queryKeyPrefix, queriedConfig, disableQueryChunking],
     staleTime: 1000 * 60 * 5,
     refetchOnWindowFocus: false,
     placeholderData: keepPreviousData, // no need to flash loading state when in live tail
+    disableQueryChunking,
   });
 
   const totalCount = useMemo(() => {

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

@@ -2,6 +2,7 @@ import React from 'react';
 import { ResponseJSON } from '@hyperdx/common-utils/dist/clickhouse';
 import { ClickhouseClient } from '@hyperdx/common-utils/dist/clickhouse/browser';
 import {
+  ChartConfigWithDateRange,
   ChartConfigWithOptDateRange,
   MetricsDataType,
 } from '@hyperdx/common-utils/dist/types';
@@ -73,33 +74,6 @@ const createMockQueryResponse = (data: any[]): ResponseJSON<any> => {
 
 describe('useChartConfig', () => {
   describe('getGranularityAlignedTimeWindows', () => {
-    it('returns [undefined] if no dateRange is provided', () => {
-      expect(
-        getGranularityAlignedTimeWindows({
-          granularity: '1 hour',
-          timestampValueExpression: 'TimestampTime',
-        } as ChartConfigWithOptDateRange),
-      ).toEqual([undefined]);
-    });
-
-    it('returns [undefined] if no granularity is provided', () => {
-      expect(
-        getGranularityAlignedTimeWindows({
-          dateRange: [new Date('2023-01-01'), new Date('2023-01-02')],
-          timestampValueExpression: 'TimestampTime',
-        } as ChartConfigWithOptDateRange),
-      ).toEqual([undefined]);
-    });
-
-    it('returns [undefined] if no timestampValueExpression is provided', () => {
-      expect(
-        getGranularityAlignedTimeWindows({
-          dateRange: [new Date('2023-01-01'), new Date('2023-01-02')],
-          granularity: '1 hour',
-        } as ChartConfigWithOptDateRange),
-      ).toEqual([undefined]);
-    });
-
     it('returns windows aligned to the granularity if the granularity is auto', () => {
       expect(
         getGranularityAlignedTimeWindows(
@@ -110,7 +84,7 @@ describe('useChartConfig', () => {
             ],
             granularity: 'auto', // will be 1 minute
             timestampValueExpression: 'TimestampTime',
-          } as ChartConfigWithOptDateRange,
+          } as ChartConfigWithDateRange & { granularity: string },
           [
             30, // 30s
             5 * 60, // 5m
@@ -152,7 +126,7 @@ describe('useChartConfig', () => {
             ],
             granularity: '1 minute',
             timestampValueExpression: 'TimestampTime',
-          } as ChartConfigWithOptDateRange,
+          } as ChartConfigWithDateRange & { granularity: string },
           [
             30, // 30s
             60, // 1m
@@ -201,7 +175,7 @@ describe('useChartConfig', () => {
             ],
             granularity: '1 minute',
             timestampValueExpression: 'TimestampTime',
-          } as ChartConfigWithOptDateRange,
+          } as ChartConfigWithDateRange & { granularity: string },
           [
             15, // 15s
           ],
@@ -235,7 +209,7 @@ describe('useChartConfig', () => {
             granularity: '1 minute',
             timestampValueExpression: 'TimestampTime',
             dateRangeEndInclusive: true,
-          } as ChartConfigWithOptDateRange,
+          } as ChartConfigWithDateRange & { granularity: string },
           [
             15 * 60, // 15m
             30 * 60, // 30m
@@ -290,7 +264,7 @@ describe('useChartConfig', () => {
             ],
             granularity: '1 minute',
             timestampValueExpression: 'TimestampTime',
-          } as ChartConfigWithOptDateRange,
+          } as ChartConfigWithDateRange & { granularity: string },
           [
             60, // 1m
           ],
@@ -885,5 +859,73 @@ describe('useChartConfig', () => {
       expect(result.current.isPending).toBe(true);
       expect(result.current.data).toBeUndefined();
     });
+
+    it('uses different query keys for the same config when one sets disableQueryChunking', async () => {
+      const config = createMockChartConfig({
+        dateRange: [
+          new Date('2025-10-01 00:00:00Z'),
+          new Date('2025-10-02 00:00:00Z'),
+        ],
+        granularity: '3 hour',
+      });
+
+      const mockResponseChunked = createMockQueryResponse([
+        {
+          'count()': '50',
+          __hdx_time_bucket: '2025-10-01T18:00:00Z',
+        },
+      ]);
+
+      const mockResponseNonChunked = createMockQueryResponse([
+        {
+          'count()': '100',
+          __hdx_time_bucket: '2025-10-01T12:00:00Z',
+        },
+      ]);
+
+      mockClickhouseClient.queryChartConfig.mockResolvedValue(
+        mockResponseChunked,
+      );
+
+      const { result: result1 } = renderHook(
+        () => useQueriedChartConfig(config),
+        {
+          wrapper,
+        },
+      );
+
+      await waitFor(() => expect(result1.current.isSuccess).toBe(true));
+      await waitFor(() => expect(result1.current.isFetching).toBe(false));
+
+      // Should have been called multiple times for chunked query
+      const chunkedCallCount =
+        mockClickhouseClient.queryChartConfig.mock.calls.length;
+      expect(chunkedCallCount).toBeGreaterThan(1);
+      expect(result1.current.data?.rows).toBeGreaterThan(1);
+
+      // Second render with same config but disableQueryChunking=true
+      mockClickhouseClient.queryChartConfig.mockResolvedValue(
+        mockResponseNonChunked,
+      );
+
+      const { result: result2 } = renderHook(
+        () => useQueriedChartConfig(config, { disableQueryChunking: true }),
+        {
+          wrapper,
+        },
+      );
+
+      await waitFor(() => expect(result2.current.isSuccess).toBe(true));
+      await waitFor(() => expect(result2.current.isFetching).toBe(false));
+
+      // Should have made a new request (not using cached chunked data)
+      expect(mockClickhouseClient.queryChartConfig).toHaveBeenCalledTimes(
+        chunkedCallCount + 1,
+      );
+      expect(result2.current.data?.rows).toBe(1);
+
+      // The original query should still have its chunked data
+      expect(result1.current.data?.rows).toBeGreaterThan(1);
+    });
   });
 });

packages/app/src/hooks/useChartConfig.tsx

@@ -65,17 +65,13 @@ const shouldUseChunking = (
   // complex metric queries.
   if (isMetricChartConfig(config)) return false;
 
-  console.log('Chunking enabled for chart config:', config);
-
   return true;
 };
 
 export const getGranularityAlignedTimeWindows = (
-  config: ChartConfigWithOptDateRange,
+  config: ChartConfigWithDateRange & { granularity: string },
   windowDurationsSeconds?: number[],
-): TimeWindow[] | [undefined] => {
-  if (!shouldUseChunking(config)) return [undefined];
-
+): TimeWindow[] => {
   const [startDate, endDate] = config.dateRange;
   const windowsUnaligned = generateTimeWindowsDescending(
     startDate,
@@ -124,9 +120,10 @@ async function* fetchDataInChunks(
   signal: AbortSignal,
   disableQueryChunking: boolean = false,
 ) {
-  const windows = disableQueryChunking
-    ? ([undefined] as const)
-    : getGranularityAlignedTimeWindows(config);
+  const windows =
+    !disableQueryChunking && shouldUseChunking(config)
+      ? getGranularityAlignedTimeWindows(config)
+      : ([undefined] as const);
 
   let query = null;
   if (IS_MTVIEWS_ENABLED) {
@@ -179,7 +176,9 @@ export function useQueriedChartConfig(
   const clickhouseClient = useClickhouseClient();
 
   const query = useQuery<TQueryFnData, ClickHouseQueryError | Error>({
-    queryKey: [config],
+    // Include disableQueryChunking in the query key to ensure that queries with the
+    // same config but different disableQueryChunking values do not share a query
+    queryKey: [config, options?.disableQueryChunking ?? false],
     queryFn: streamedQuery({
       streamFn: context =>
         fetchDataInChunks(