perf: add prelimit CTE to getMapKeys query + store clickhouse settings in shared cache (#1201)

Optimize the query performance of the getMapKeys method to prevent excessive resource usage in ClickHouse, even when max_rows_to_read is specified.

Ref: HDX-2411
Ref: HDX-2431

Author: wrn14897

.changeset/strange-boats-nail.md

@@ -0,0 +1,6 @@
+---
+"@hyperdx/common-utils": patch
+"@hyperdx/app": patch
+---
+
+perf: add prelimit CTE to getMapKeys query + store clickhouse settings in shared cache

packages/app/src/components/DBSearchPageFilters.tsx

@@ -33,8 +33,8 @@ import {
   useJsonColumns,
   useTableMetadata,
 } from '@/hooks/useMetadata';
+import { useMetadataWithSettings } from '@/hooks/useMetadata';
 import useResizable from '@/hooks/useResizable';
-import { getMetadata } from '@/metadata';
 import { FilterStateHook, usePinnedFilters } from '@/searchFilters';
 import { useSource } from '@/source';
 import { mergePath } from '@/utils';
@@ -615,6 +615,7 @@ const DBSearchPageFiltersComponent = ({
     keys: keysToFetch,
   });
 
+  const metadata = useMetadataWithSettings();
   const [extraFacets, setExtraFacets] = useState<Record<string, string[]>>({});
   const [loadMoreLoadingKeys, setLoadMoreLoadingKeys] = useState<Set<string>>(
     new Set(),
@@ -623,7 +624,6 @@ const DBSearchPageFiltersComponent = ({
     async (key: string) => {
       setLoadMoreLoadingKeys(prev => new Set(prev).add(key));
       try {
-        const metadata = getMetadata();
         const newKeyVals = await metadata.getKeyValues({
           chartConfig: {
             ...chartConfig,

packages/app/src/hooks/useMetadata.tsx

@@ -1,3 +1,4 @@
+import { useEffect, useRef } from 'react';
 import objectHash from 'object-hash';
 import {
   ColumnMeta,
@@ -20,6 +21,25 @@ import api from '@/api';
 import { getMetadata } from '@/metadata';
 import { toArray } from '@/utils';
 
+// Hook to get metadata with proper settings applied
+// TODO: replace all getMetadata calls with useMetadataWithSettings
+export function useMetadataWithSettings() {
+  const metadata = getMetadata();
+  const { data: me } = api.useMe();
+  const settingsApplied = useRef(false);
+
+  useEffect(() => {
+    if (me?.team?.metadataMaxRowsToRead && !settingsApplied.current) {
+      metadata.setClickHouseSettings({
+        max_rows_to_read: me.team.metadataMaxRowsToRead,
+      });
+      settingsApplied.current = true;
+    }
+  }, [me?.team?.metadataMaxRowsToRead, metadata]);
+
+  return metadata;
+}
+
 export function useColumns(
   {
     databaseName,
@@ -32,10 +52,10 @@ export function useColumns(
   },
   options?: Partial<UseQueryOptions<ColumnMeta[]>>,
 ) {
+  const metadata = useMetadataWithSettings();
   return useQuery<ColumnMeta[]>({
     queryKey: ['useMetadata.useColumns', { databaseName, tableName }],
     queryFn: async () => {
-      const metadata = getMetadata();
       return metadata.getColumns({
         databaseName,
         tableName,
@@ -51,10 +71,10 @@ export function useJsonColumns(
   { databaseName, tableName, connectionId }: TableConnection,
   options?: Partial<UseQueryOptions<string[]>>,
 ) {
+  const metadata = useMetadataWithSettings();
   return useQuery<string[]>({
     queryKey: ['useMetadata.useJsonColumns', { databaseName, tableName }],
     queryFn: async () => {
-      const metadata = getMetadata();
       const columns = await metadata.getColumns({
         databaseName,
         tableName,
@@ -78,7 +98,7 @@ export function useAllFields(
   const tableConnections = Array.isArray(_tableConnections)
     ? _tableConnections
     : [_tableConnections];
-  const metadata = getMetadata();
+  const metadata = useMetadataWithSettings();
   const { data: me, isFetched } = api.useMe();
   return useQuery<Field[]>({
     queryKey: [
@@ -91,13 +111,6 @@ export function useAllFields(
         return [];
       }
 
-      // TODO: set the settings at the top level so that it doesn't have to be set for each useQuery
-      if (team?.metadataMaxRowsToRead) {
-        metadata.setClickHouseSettings({
-          max_rows_to_read: team.metadataMaxRowsToRead,
-        });
-      }
-
       const fields2d = await Promise.all(
         tableConnections.map(tc => metadata.getAllFields(tc)),
       );
@@ -129,7 +142,7 @@ export function useTableMetadata(
   },
   options?: Omit<UseQueryOptions<any, Error>, 'queryKey'>,
 ) {
-  const metadata = getMetadata();
+  const metadata = useMetadataWithSettings();
   return useQuery<TableMetadata>({
     queryKey: ['useMetadata.useTableMetadata', { databaseName, tableName }],
     queryFn: async () => {
@@ -159,9 +172,8 @@ export function useGetKeyValues(
   },
   options?: Omit<UseQueryOptions<any, Error>, 'queryKey'>,
 ) {
-  const metadata = getMetadata();
+  const metadata = useMetadataWithSettings();
   const chartConfigsArr = toArray(chartConfigs);
-  const { data: me, isFetched } = api.useMe();
   return useQuery<{ key: string; value: string[] }[]>({
     queryKey: [
       'useMetadata.useGetKeyValues',
@@ -170,14 +182,6 @@ export function useGetKeyValues(
       disableRowLimit,
     ],
     queryFn: async () => {
-      const team = me?.team;
-
-      // TODO: set the settings at the top level so that it doesn't have to be set for each useQuery
-      if (team?.metadataMaxRowsToRead) {
-        metadata.setClickHouseSettings({
-          max_rows_to_read: team.metadataMaxRowsToRead,
-        });
-      }
       return (
         await Promise.all(
           chartConfigsArr.map(chartConfig =>
@@ -192,7 +196,7 @@ export function useGetKeyValues(
       ).flatMap(v => v);
     },
     staleTime: 1000 * 60 * 5, // Cache every 5 min
-    enabled: !!keys.length && isFetched,
+    enabled: !!keys.length,
     placeholderData: keepPreviousData,
     ...options,
   });

packages/common-utils/src/__tests__/metadata.test.ts

@@ -251,6 +251,22 @@ describe('Metadata', () => {
             ],
           }),
       });
+
+      // Mock getColumns to return columns including materialized ones
+      mockCache.getOrFetch.mockImplementation((key, queryFn) => {
+        if (key.includes('.columns')) {
+          return Promise.resolve([
+            { name: 'regular_column', type: 'String', default_type: '' },
+            {
+              name: 'materialized_column',
+              type: 'String',
+              default_type: 'MATERIALIZED',
+            },
+            { name: 'default_column', type: 'String', default_type: 'DEFAULT' },
+          ]);
+        }
+        return queryFn();
+      });
     });
 
     it('should apply row limit when disableRowLimit is false', async () => {
@@ -336,5 +352,72 @@ describe('Metadata', () => {
 
       expect(result).toEqual([{ key: 'column1', value: ['value1', 'value2'] }]);
     });
+
+    it('should include materialized fields when selecting all columns', async () => {
+      const renderChartConfigSpy = jest.spyOn(
+        renderChartConfigModule,
+        'renderChartConfig',
+      );
+
+      await metadata.getKeyValues({
+        chartConfig: mockChartConfig,
+        keys: ['column1'],
+        limit: 10,
+      });
+
+      // Verify that renderChartConfig was called with the expanded select list
+      // that includes all columns by name (including materialized ones)
+      expect(renderChartConfigSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          with: [
+            expect.objectContaining({
+              name: 'sampledData',
+              chartConfig: expect.objectContaining({
+                // Should expand to all column names instead of using '*'
+                select:
+                  '`regular_column`, `materialized_column`, `default_column`',
+              }),
+            }),
+          ],
+        }),
+        metadata,
+      );
+    });
+
+    it('should fallback to * when no columns are found', async () => {
+      // Mock getColumns to return empty array
+      mockCache.getOrFetch.mockImplementation((key, queryFn) => {
+        if (key.includes('.columns')) {
+          return Promise.resolve([]);
+        }
+        return queryFn();
+      });
+
+      const renderChartConfigSpy = jest.spyOn(
+        renderChartConfigModule,
+        'renderChartConfig',
+      );
+
+      await metadata.getKeyValues({
+        chartConfig: mockChartConfig,
+        keys: ['column1'],
+        limit: 10,
+      });
+
+      // Should fallback to '*' when no columns are found
+      expect(renderChartConfigSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          with: [
+            expect.objectContaining({
+              name: 'sampledData',
+              chartConfig: expect.objectContaining({
+                select: '*',
+              }),
+            }),
+          ],
+        }),
+        metadata,
+      );
+    });
   });
 });