feat: Separate useChunkedQueriedChartConfig from useQueriedChartConfig

pulpdrew · pulpdrew · commit 3e899c8e6eed · 2025-10-06T15:17:28.000-04:00
diff --git a/packages/app/src/hooks/__tests__/useChartConfig.test.tsx b/packages/app/src/hooks/__tests__/useChartConfig.test.tsx
@@ -9,7 +9,7 @@ import { useClickhouseClient } from '@/clickhouse';
 
 import {
   getGranularityAlignedTimeWindows,
-  useQueriedChartConfig,
+  useChunkedQueriedChartConfig,
 } from '../useChartConfig';
 
 // Mock the clickhouse module
@@ -97,6 +97,48 @@ describe('useChartConfig', () => {
       ).toEqual([undefined]);
     });
 
+    it('returns windows aligned to the granularity if the granularity is auto', () => {
+      expect(
+        getGranularityAlignedTimeWindows(
+          {
+            dateRange: [
+              new Date('2023-01-10 00:00:00'),
+              new Date('2023-01-10 01:00:00'),
+            ],
+            granularity: 'auto', // will be 1 minute
+            timestampValueExpression: 'TimestampTime',
+          } as ChartConfigWithOptDateRange,
+          [
+            30, // 30s
+            5 * 60, // 5m
+            60 * 60, // 1hr
+          ],
+        ),
+      ).toEqual([
+        {
+          dateRange: [
+            new Date('2023-01-10 00:59:00'), // Aligned to minute, the auto-inferred granularity
+            new Date('2023-01-10 01:00:00'),
+          ],
+          dateRangeEndInclusive: undefined,
+        },
+        {
+          dateRange: [
+            new Date('2023-01-10 00:54:00'),
+            new Date('2023-01-10 00:59:00'),
+          ],
+          dateRangeEndInclusive: false,
+        },
+        {
+          dateRange: [
+            new Date('2023-01-10 00:00:00'),
+            new Date('2023-01-10 00:54:00'),
+          ],
+          dateRangeEndInclusive: false,
+        },
+      ]);
+    });
+
     it('returns windows aligned to the granularity if the granularity is larger than the window size', () => {
       expect(
         getGranularityAlignedTimeWindows(
@@ -269,7 +311,7 @@ describe('useChartConfig', () => {
     });
   });
 
-  describe('useQueriedChartConfig', () => {
+  describe('useChunkedQueriedChartConfig', () => {
     let queryClient: QueryClient;
     let wrapper: React.ComponentType<{ children: any }>;
     let mockClickhouseClient: jest.Mocked<ClickhouseClient>;
@@ -319,9 +361,12 @@ describe('useChartConfig', () => {
 
       mockClickhouseClient.queryChartConfig.mockResolvedValue(mockResponse);
 
-      const { result } = renderHook(() => useQueriedChartConfig(config), {
-        wrapper,
-      });
+      const { result } = renderHook(
+        () => useChunkedQueriedChartConfig(config),
+        {
+          wrapper,
+        },
+      );
 
       await waitFor(() => expect(result.current.isSuccess).toBe(true));
       await waitFor(() => expect(result.current.isFetching).toBe(false));
@@ -364,9 +409,12 @@ describe('useChartConfig', () => {
 
       mockClickhouseClient.queryChartConfig.mockResolvedValue(mockResponse);
 
-      const { result } = renderHook(() => useQueriedChartConfig(config), {
-        wrapper,
-      });
+      const { result } = renderHook(
+        () => useChunkedQueriedChartConfig(config),
+        {
+          wrapper,
+        },
+      );
 
       await waitFor(() => expect(result.current.isSuccess).toBe(true));
       await waitFor(() => expect(result.current.isFetching).toBe(false));
@@ -413,55 +461,8 @@ describe('useChartConfig', () => {
 
       mockClickhouseClient.queryChartConfig.mockResolvedValue(mockResponse);
 
-      const { result } = renderHook(() => useQueriedChartConfig(config), {
-        wrapper,
-      });
-
-      await waitFor(() => expect(result.current.isSuccess).toBe(true));
-      await waitFor(() => expect(result.current.isFetching).toBe(false));
-
-      // Should only be called once since chunking is disabled without timestampValueExpression
-      expect(mockClickhouseClient.queryChartConfig).toHaveBeenCalledTimes(1);
-      expect(mockClickhouseClient.queryChartConfig).toHaveBeenCalledWith({
-        config,
-        metadata: expect.any(Object),
-        opts: {
-          abort_signal: expect.any(AbortSignal),
-        },
-      });
-      expect(result.current.data).toEqual({
-        data: mockResponse.data,
-        meta: mockResponse.meta,
-        rows: mockResponse.rows,
-      });
-    });
-
-    it('fetches data without chunking when disableQueryChunking is true', async () => {
-      const config = createMockChartConfig({
-        dateRange: [
-          new Date('2025-10-01 00:00:00Z'),
-          new Date('2025-10-02 00:00:00Z'),
-        ],
-        granularity: '1 hour',
-      });
-
-      const mockResponse = createMockQueryResponse([
-        {
-          'count()': '71',
-          SeverityText: 'info',
-          __hdx_time_bucket: '2025-10-01T00:00:00Z',
-        },
-        {
-          'count()': '73',
-          SeverityText: 'info',
-          __hdx_time_bucket: '2025-10-02T00:00:00Z',
-        },
-      ]);
-
-      mockClickhouseClient.queryChartConfig.mockResolvedValue(mockResponse);
-
       const { result } = renderHook(
-        () => useQueriedChartConfig(config, { disableQueryChunking: true }),
+        () => useChunkedQueriedChartConfig(config),
         {
           wrapper,
         },
@@ -470,7 +471,7 @@ describe('useChartConfig', () => {
       await waitFor(() => expect(result.current.isSuccess).toBe(true));
       await waitFor(() => expect(result.current.isFetching).toBe(false));
 
-      // Should only be called once since chunking is explicitly disabled
+      // Should only be called once since chunking is disabled without timestampValueExpression
       expect(mockClickhouseClient.queryChartConfig).toHaveBeenCalledTimes(1);
       expect(mockClickhouseClient.queryChartConfig).toHaveBeenCalledWith({
         config,
@@ -529,9 +530,12 @@ describe('useChartConfig', () => {
         .mockResolvedValueOnce(mockResponse2)
         .mockResolvedValueOnce(mockResponse3);
 
-      const { result } = renderHook(() => useQueriedChartConfig(config), {
-        wrapper,
-      });
+      const { result } = renderHook(
+        () => useChunkedQueriedChartConfig(config),
+        {
+          wrapper,
+        },
+      );
 
       await waitFor(() => expect(result.current.isSuccess).toBe(true));
       await waitFor(() => expect(result.current.isFetching).toBe(false));
@@ -620,9 +624,12 @@ describe('useChartConfig', () => {
         .mockResolvedValueOnce(mockResponse2)
         .mockResolvedValueOnce(mockResponse3);
 
-      const { result } = renderHook(() => useQueriedChartConfig(config), {
-        wrapper,
-      });
+      const { result } = renderHook(
+        () => useChunkedQueriedChartConfig(config),
+        {
+          wrapper,
+        },
+      );
 
       await waitFor(() => expect(result.current.isSuccess).toBe(true));
       await waitFor(() => expect(result.current.isPending).toBe(false));
@@ -678,9 +685,12 @@ describe('useChartConfig', () => {
         mockResponse1Promise,
       );
 
-      const { result } = renderHook(() => useQueriedChartConfig(config), {
-        wrapper,
-      });
+      const { result } = renderHook(
+        () => useChunkedQueriedChartConfig(config),
+        {
+          wrapper,
+        },
+      );
 
       // Should be in loading state before first chunk
       expect(result.current.isLoading).toBe(true);
@@ -721,7 +731,11 @@ describe('useChartConfig', () => {
       });
 
       const { result } = renderHook(
-        () => useQueriedChartConfig(config, { onError, retry: false }),
+        () =>
+          useChunkedQueriedChartConfig(config, {
+            onError,
+            retry: false,
+          }),
         {
           wrapper,
         },
@@ -752,7 +766,10 @@ describe('useChartConfig', () => {
       mockClickhouseClient.queryChartConfig.mockResolvedValue(mockResponse);
 
       const { result } = renderHook(
-        () => useQueriedChartConfig(config, { enabled: false }),
+        () =>
+          useChunkedQueriedChartConfig(config, {
+            enabled: false,
+          }),
         {
           wrapper,
         },
diff --git a/packages/app/src/hooks/useChartConfig.tsx b/packages/app/src/hooks/useChartConfig.tsx
@@ -30,13 +30,6 @@ import { generateTimeWindowsDescending } from '@/utils/searchWindows';
 
 interface AdditionalUseQueriedChartConfigOptions {
   onError?: (error: Error | ClickHouseQueryError) => void;
-  /**
-   * By default, queries with large date ranges are split into multiple smaller queries to
-   * avoid overloading the ClickHouse server and running into timeouts. In some cases, such
-   * as when data is being sampled across the entire range, this chunking is not desirable
-   * and can be disabled.
-   */
-  disableQueryChunking?: boolean;
 }
 
 type TimeWindow = {
@@ -64,7 +57,7 @@ export const getGranularityAlignedTimeWindows = (
   );
 
   const granularity =
-    config.granularity === 'auto' // TODO test this
+    config.granularity === 'auto'
       ? convertDateRangeToGranularityString(
           config.dateRange,
           DEFAULT_AUTO_GRANULARITY_MAX_BUCKETS,
@@ -102,24 +95,8 @@ async function* fetchDataInChunks(
   config: ChartConfigWithOptDateRange,
   clickhouseClient: ClickhouseClient,
   signal: AbortSignal,
-  disableQueryChunking: boolean = false,
 ) {
-  const windows = disableQueryChunking
-    ? ([undefined] as const)
-    : getGranularityAlignedTimeWindows(config);
-
-  let query = null;
-  if (IS_MTVIEWS_ENABLED) {
-    const { dataTableDDL, mtViewDDL, renderMTViewConfig } =
-      await buildMTViewSelectQuery(config);
-    // TODO: show the DDLs in the UI so users can run commands manually
-    // eslint-disable-next-line no-console
-    console.log('dataTableDDL:', dataTableDDL);
-    // eslint-disable-next-line no-console
-    console.log('mtViewDDL:', mtViewDDL);
-    query = await renderMTViewConfig();
-  }
-
+  const windows = getGranularityAlignedTimeWindows(config);
   for (const window of windows) {
     const windowedConfig = {
       ...config,
@@ -141,18 +118,16 @@ async function* fetchDataInChunks(
 /**
  * A hook providing data queried based on the provided chart config.
  *
- * If all of the following are true, the query will be chunked into multiple smaller queries:
- * - The config includes a date range
- * - The config includes a granularity
- * - `options.disableQueryChunking` is falsy
+ * The query will be chunked into multiple smaller queries if the config
+ * includes a dateRange, granularity, timestampValueExpression.
  *
  * For chunked queries, note the following:
- * - The config's limit, if provided, is applied to each chunk, so the total number
+ * - `config.limit`, if provided, is applied to each chunk, so the total number
  *   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`.
  */
-export function useQueriedChartConfig(
+export function useChunkedQueriedChartConfig(
   config: ChartConfigWithOptDateRange,
   options?: Partial<UseQueryOptions<ResponseJSON<any>>> &
     AdditionalUseQueriedChartConfigOptions,
@@ -163,12 +138,7 @@ export function useQueriedChartConfig(
     queryKey: [config],
     queryFn: streamedQuery({
       streamFn: context =>
-        fetchDataInChunks(
-          config,
-          clickhouseClient,
-          context.signal,
-          options?.disableQueryChunking,
-        ),
+        fetchDataInChunks(config, clickhouseClient, context.signal),
       /**
        * This mode ensures that data remains in the cache until the next full streamed result is available.
        * By default, the cache would be cleared before new data starts arriving, which results in the query briefly
@@ -197,6 +167,45 @@ export function useQueriedChartConfig(
   return query;
 }
 
+export function useQueriedChartConfig(
+  config: ChartConfigWithOptDateRange,
+  options?: Partial<UseQueryOptions<ResponseJSON<any>>> &
+    AdditionalUseQueriedChartConfigOptions,
+) {
+  const clickhouseClient = useClickhouseClient();
+  const query = useQuery<ResponseJSON<any>, ClickHouseQueryError | Error>({
+    queryKey: [config],
+    queryFn: async ({ signal }) => {
+      let query = null;
+      if (IS_MTVIEWS_ENABLED) {
+        const { dataTableDDL, mtViewDDL, renderMTViewConfig } =
+          await buildMTViewSelectQuery(config);
+        // TODO: show the DDLs in the UI so users can run commands manually
+        // eslint-disable-next-line no-console
+        console.log('dataTableDDL:', dataTableDDL);
+        // eslint-disable-next-line no-console
+        console.log('mtViewDDL:', mtViewDDL);
+        query = await renderMTViewConfig();
+      }
+
+      return clickhouseClient.queryChartConfig({
+        config,
+        metadata: getMetadata(),
+        opts: {
+          abort_signal: signal,
+        },
+      });
+    },
+    retry: 1,
+    refetchOnWindowFocus: false,
+    ...options,
+  });
+  if (query.isError && options?.onError) {
+    options.onError(query.error);
+  }
+  return query;
+}
+
 export function useRenderedSqlChartConfig(
   config: ChartConfigWithOptDateRange,
   options?: UseQueryOptions<string>,
diff --git a/packages/app/src/hooks/usePatterns.tsx b/packages/app/src/hooks/usePatterns.tsx
@@ -143,11 +143,7 @@ function usePatterns({
 
   const { data: sampleRows } = useQueriedChartConfig(
     configWithPrimaryAndPartitionKey ?? config, // `config` satisfying type, never used due to `enabled` check
-    {
-      enabled: configWithPrimaryAndPartitionKey != null && enabled,
-      // Disable chunking so that we get a uniform sample across the entire date range
-      disableQueryChunking: true,
-    },
+    { enabled: configWithPrimaryAndPartitionKey != null && enabled },
   );
 
   const { data: pyodide, isLoading: isLoadingPyodide } = usePyodide({
