diff --git a/.changeset/strange-boats-nail.md b/.changeset/strange-boats-nail.md new file mode 100644 index 0000000000..7f2e62534e --- /dev/null +++ b/.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 diff --git a/packages/app/src/components/DBSearchPageFilters.tsx b/packages/app/src/components/DBSearchPageFilters.tsx index 2c6489e5e6..ce69b52fa1 100644 --- a/packages/app/src/components/DBSearchPageFilters.tsx +++ b/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>({}); const [loadMoreLoadingKeys, setLoadMoreLoadingKeys] = useState>( 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, diff --git a/packages/app/src/hooks/useMetadata.tsx b/packages/app/src/hooks/useMetadata.tsx index f4134a8af7..506a04b8a8 100644 --- a/packages/app/src/hooks/useMetadata.tsx +++ b/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>, ) { + const metadata = useMetadataWithSettings(); return useQuery({ 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>, ) { + const metadata = useMetadataWithSettings(); return useQuery({ 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({ 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, 'queryKey'>, ) { - const metadata = getMetadata(); + const metadata = useMetadataWithSettings(); return useQuery({ queryKey: ['useMetadata.useTableMetadata', { databaseName, tableName }], queryFn: async () => { @@ -159,9 +172,8 @@ export function useGetKeyValues( }, options?: Omit, '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, }); diff --git a/packages/common-utils/src/__tests__/metadata.test.ts b/packages/common-utils/src/__tests__/metadata.test.ts index 11668dd7bd..b1b7902409 100644 --- a/packages/common-utils/src/__tests__/metadata.test.ts +++ b/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, + ); + }); }); }); diff --git a/packages/common-utils/src/metadata.ts b/packages/common-utils/src/metadata.ts index a0d6f3f90c..8cd3206d26 100644 --- a/packages/common-utils/src/metadata.ts +++ b/packages/common-utils/src/metadata.ts @@ -94,7 +94,6 @@ export type TableMetadata = { export class Metadata { private readonly clickhouseClient: BaseClickhouseClient; private readonly cache: MetadataCache; - private readonly clickhouseSettings: ClickHouseSettings; constructor( clickhouseClient: BaseClickhouseClient, @@ -103,11 +102,19 @@ export class Metadata { ) { this.clickhouseClient = clickhouseClient; this.cache = cache; - this.clickhouseSettings = settings ?? {}; + if (settings) { + this.cache.set('clickhouse-settings', settings); + } + } + + getClickHouseSettings(): ClickHouseSettings { + return this.cache.get('clickhouse-settings') ?? {}; } setClickHouseSettings(settings: ClickHouseSettings) { - Object.assign(this.clickhouseSettings, settings); + const currentSettings = this.getClickHouseSettings(); + const updatedSettings = { ...currentSettings, ...settings }; + this.cache.set('clickhouse-settings', updatedSettings); } private async queryTableMetadata({ @@ -128,7 +135,7 @@ export class Metadata { connectionId, query: sql.sql, query_params: sql.params, - clickhouse_settings: this.clickhouseSettings, + clickhouse_settings: this.getClickHouseSettings(), }) .then(res => res.json()); return json.data[0]; @@ -153,7 +160,7 @@ export class Metadata { query: sql.sql, query_params: sql.params, connectionId, - clickhouse_settings: this.clickhouseSettings, + clickhouse_settings: this.getClickHouseSettings(), }) .then(res => res.json()) .then(d => d.data); @@ -264,18 +271,39 @@ export class Metadata { : ''; let sql: ChSql; if (strategy === 'groupUniqArrayArray') { - sql = chSql`SELECT groupUniqArrayArray(${{ Int32: maxKeys }})(${{ - Identifier: column, - }}) as keysArr - FROM ${tableExpr({ database: databaseName, table: tableName })} ${where}`; + sql = chSql` + WITH sampledKeys as ( + SELECT ${{ + Identifier: column, + }}.keys AS keys + FROM ${tableExpr({ database: databaseName, table: tableName })} ${where} + LIMIT ${{ + Int32: this.getClickHouseSettings().max_rows_to_read + ? Number(this.getClickHouseSettings().max_rows_to_read) + : DEFAULT_METADATA_MAX_ROWS_TO_READ, + }} + ) + SELECT groupUniqArrayArray(${{ Int32: maxKeys }})(keys) as keysArr + FROM sampledKeys`; } else { - sql = chSql`SELECT DISTINCT lowCardinalityKeys(arrayJoin(${{ - Identifier: column, - }}.keys)) as key - FROM ${tableExpr({ database: databaseName, table: tableName })} ${where} - LIMIT ${{ - Int32: maxKeys, - }}`; + sql = chSql` + WITH sampledKeys as ( + SELECT ${{ + Identifier: column, + }}.keys AS keysArr + FROM ${tableExpr({ database: databaseName, table: tableName })} ${where} + LIMIT ${{ + Int32: this.getClickHouseSettings().max_rows_to_read + ? Number(this.getClickHouseSettings().max_rows_to_read) + : DEFAULT_METADATA_MAX_ROWS_TO_READ, + }} + ) + SELECT DISTINCT lowCardinalityKeys(arrayJoin(keysArr)) as key + FROM sampledKeys + LIMIT ${{ + Int32: maxKeys, + }} + `; } return this.cache.getOrFetch(cacheKey, async () => { @@ -285,9 +313,12 @@ export class Metadata { query_params: sql.params, connectionId, clickhouse_settings: { - max_rows_to_read: String(DEFAULT_METADATA_MAX_ROWS_TO_READ), + max_rows_to_read: String( + this.getClickHouseSettings().max_rows_to_read ?? + DEFAULT_METADATA_MAX_ROWS_TO_READ, + ), read_overflow_mode: 'break', - ...this.clickhouseSettings, + ...this.getClickHouseSettings(), }, }) .then(res => res.json>()) @@ -342,9 +373,12 @@ export class Metadata { query_params: sql.params, connectionId, clickhouse_settings: { - max_rows_to_read: String(DEFAULT_METADATA_MAX_ROWS_TO_READ), + max_rows_to_read: String( + this.getClickHouseSettings().max_rows_to_read ?? + DEFAULT_METADATA_MAX_ROWS_TO_READ, + ), read_overflow_mode: 'break', - ...this.clickhouseSettings, + ...this.getClickHouseSettings(), }, }) .then(res => res.json<{ pathMap: Record }>()) @@ -422,9 +456,12 @@ export class Metadata { query_params: sql.params, connectionId, clickhouse_settings: { - max_rows_to_read: String(DEFAULT_METADATA_MAX_ROWS_TO_READ), + max_rows_to_read: String( + this.getClickHouseSettings().max_rows_to_read ?? + DEFAULT_METADATA_MAX_ROWS_TO_READ, + ), read_overflow_mode: 'break', - ...this.clickhouseSettings, + ...this.getClickHouseSettings(), }, }) .then(res => res.json>()) @@ -543,12 +580,43 @@ export class Metadata { return this.cache.getOrFetch( `${chartConfig.from.databaseName}.${chartConfig.from.tableName}.${keys.join(',')}.${chartConfig.dateRange.toString()}.${disableRowLimit}.values`, async () => { + // Get all columns including materialized ones + const columns = await this.getColumns({ + databaseName: chartConfig.from.databaseName, + tableName: chartConfig.from.tableName, + connectionId: chartConfig.connection, + }); + + // Build select expression that includes all columns by name + // This ensures materialized columns are included + const selectExpr = + columns.map(col => `\`${col.name}\``).join(', ') || '*'; + const sql = await renderChartConfig( { - ...chartConfig, + with: [ + { + name: 'sampledData', + chartConfig: { + ...chartConfig, + select: selectExpr, + limit: { + limit: !disableRowLimit + ? this.getClickHouseSettings().max_rows_to_read + ? Number(this.getClickHouseSettings().max_rows_to_read) + : DEFAULT_METADATA_MAX_ROWS_TO_READ + : undefined, + }, + }, + isSubquery: true, + }, + ], select: keys .map((k, i) => `groupUniqArray(${limit})(${k}) AS param${i}`) .join(', '), + connection: chartConfig.connection, + from: { databaseName: '', tableName: 'sampledData' }, + where: '', }, this, ); @@ -560,9 +628,12 @@ export class Metadata { connectionId: chartConfig.connection, clickhouse_settings: !disableRowLimit ? { - max_rows_to_read: String(DEFAULT_METADATA_MAX_ROWS_TO_READ), + max_rows_to_read: String( + this.getClickHouseSettings().max_rows_to_read ?? + DEFAULT_METADATA_MAX_ROWS_TO_READ, + ), read_overflow_mode: 'break', - ...this.clickhouseSettings, + ...this.getClickHouseSettings(), } : undefined, })