Team query timeout setting (#1072)

Closes HDX-2023

# Behavior
Adds `queryTimeout` to the team settings page. If the setting is greater than `0`, a `max_execution_time` setting is added to all queries executed by the `ClickhouseClient`. The timeout is in seconds, with `60` seconds being the default. `0` would be unlimited.

# Images

Settings page:
<img width="598" height="501" alt="settings page" src="https://github.com/user-attachments/assets/5483d5a7-c1c2-4bb5-a0d9-e23fa06cc5ec" />

Network log:
<img width="542" height="100" alt="network log" src="https://github.com/user-attachments/assets/08ed2ea1-4038-4c67-b493-ef591a226b59" />

.changeset/afraid-rocks-smile.md

@@ -0,0 +1,7 @@
+---
+"@hyperdx/common-utils": minor
+"@hyperdx/api": minor
+"@hyperdx/app": minor
+---
+
+added team level queryTimeout to ClickHouse client

packages/api/src/models/team.ts

@@ -6,6 +6,7 @@ type ObjectId = mongoose.Types.ObjectId;
 export type TeamCHSettings = {
   metadataMaxRowsToRead?: number;
   searchRowLimit?: number;
+  queryTimeout?: number;
   fieldMetadataDisabled?: boolean;
 };
 
@@ -46,6 +47,7 @@ export default mongoose.model<ITeam>(
       // CH Client Settings
       metadataMaxRowsToRead: Number,
       searchRowLimit: Number,
+      queryTimeout: Number,
       fieldMetadataDisabled: Boolean,
     },
     {

packages/api/src/routers/api/team.ts

@@ -94,6 +94,7 @@ router.patch(
     body: z.object({
       fieldMetadataDisabled: z.boolean().optional(),
       searchRowLimit: z.number().optional(),
+      queryTimeout: z.number().optional(),
       metadataMaxRowsToRead: z.number().optional(),
     }),
   }),
@@ -104,11 +105,16 @@ router.patch(
         throw new Error(`User ${req.user?._id} not associated with a team`);
       }
 
-      const { fieldMetadataDisabled, metadataMaxRowsToRead, searchRowLimit } =
-        req.body;
+      const {
+        fieldMetadataDisabled,
+        metadataMaxRowsToRead,
+        searchRowLimit,
+        queryTimeout,
+      } = req.body;
 
       const settings = {
         ...(searchRowLimit !== undefined && { searchRowLimit }),
+        ...(queryTimeout !== undefined && { queryTimeout }),
         ...(fieldMetadataDisabled !== undefined && { fieldMetadataDisabled }),
         ...(metadataMaxRowsToRead !== undefined && { metadataMaxRowsToRead }),
       };
@@ -123,6 +129,9 @@ router.patch(
         ...(searchRowLimit !== undefined && {
           searchRowLimit: team?.searchRowLimit,
         }),
+        ...(queryTimeout !== undefined && {
+          queryTimeout: team?.queryTimeout,
+        }),
         ...(fieldMetadataDisabled !== undefined && {
           fieldMetadataDisabled: team?.fieldMetadataDisabled,
         }),

packages/app/src/BenchmarkPage.tsx

@@ -21,7 +21,7 @@ import {
 } from '@mantine/core';
 import { useQuery, UseQueryOptions } from '@tanstack/react-query';
 
-import { getClickhouseClient } from '@/clickhouse';
+import { useClickhouseClient } from '@/clickhouse';
 
 import { ConnectionSelectControlled } from './components/ConnectionSelect';
 import DBTableChart from './components/DBTableChart';
@@ -38,7 +38,7 @@ function useBenchmarkQueryIds({
   iterations?: number;
 }) {
   const enabled = queries.length > 0 && connections.length > 0;
-  const clickhouseClient = getClickhouseClient();
+  const clickhouseClient = useClickhouseClient();
 
   return useQuery({
     enabled,
@@ -95,7 +95,7 @@ function useEstimates(
   },
   options: Omit<UseQueryOptions<any>, 'queryKey' | 'queryFn'> = {},
 ) {
-  const clickhouseClient = getClickhouseClient();
+  const clickhouseClient = useClickhouseClient();
   return useQuery({
     queryKey: ['estimate', queries, connections],
     queryFn: async () => {
@@ -125,7 +125,7 @@ function useIndexes(
   },
   options: Omit<UseQueryOptions<any>, 'queryKey' | 'queryFn'> = {},
 ) {
-  const clickhouseClient = getClickhouseClient();
+  const clickhouseClient = useClickhouseClient();
   return useQuery({
     queryKey: ['indexes', queries, connections],
     queryFn: async () => {

packages/app/src/TeamPage.tsx

@@ -43,7 +43,7 @@ import { IS_LOCAL_MODE } from '@/config';
 import { PageHeader } from './components/PageHeader';
 import api from './api';
 import { useConnections } from './connection';
-import { DEFAULT_SEARCH_ROW_LIMIT } from './defaults';
+import { DEFAULT_QUERY_TIMEOUT, DEFAULT_SEARCH_ROW_LIMIT } from './defaults';
 import { withAppNav } from './layout';
 import { useSources } from './source';
 import { useConfirm } from './useConfirm';
@@ -1069,6 +1069,7 @@ type ClickhouseSettingType = 'number' | 'boolean';
 interface ClickhouseSettingFormProps {
   settingKey:
     | 'searchRowLimit'
+    | 'queryTimeout'
     | 'metadataMaxRowsToRead'
     | 'fieldMetadataDisabled';
   label: string;
@@ -1078,7 +1079,7 @@ interface ClickhouseSettingFormProps {
   placeholder?: string;
   min?: number;
   max?: number;
-  displayValue?: (value: any) => string;
+  displayValue?: (value: any, defaultValue?: any) => string;
   options?: string[]; // For boolean settings displayed as select
 }
 
@@ -1237,7 +1238,7 @@ function ClickhouseSettingForm({
         <Group>
           <Text className="text-white">
             {displayValue
-              ? displayValue(currentValue)
+              ? displayValue(currentValue, defaultValue)
               : currentValue?.toString() || 'Not set'}
           </Text>
           {hasAdminAccess && (
@@ -1257,6 +1258,14 @@ function ClickhouseSettingForm({
 }
 
 function TeamQueryConfigSection() {
+  const displayValueWithUnit =
+    (unit: string) => (value: any, defaultValue?: any) =>
+      value === undefined || value === defaultValue
+        ? `${defaultValue.toLocaleString()} ${unit} (System Default)`
+        : value === 0
+          ? 'Unlimited'
+          : `${value.toLocaleString()} ${unit}`;
+
   return (
     <Box id="team_name">
       <Text size="md" c="gray.4">
@@ -1271,10 +1280,20 @@ function TeamQueryConfigSection() {
             tooltip="The number of rows per query for the Search page or search dashboard tiles"
             type="number"
             defaultValue={DEFAULT_SEARCH_ROW_LIMIT}
-            placeholder={`Enter value (default: ${DEFAULT_SEARCH_ROW_LIMIT})`}
+            placeholder={`default = ${DEFAULT_SEARCH_ROW_LIMIT}, 0 = unlimited`}
             min={1}
             max={100000}
-            displayValue={value => value ?? 'System Default'}
+            displayValue={displayValueWithUnit('rows')}
+          />
+          <ClickhouseSettingForm
+            settingKey="queryTimeout"
+            label="Query Timeout (seconds)"
+            tooltip="Sets the max execution time of a query in seconds."
+            type="number"
+            defaultValue={DEFAULT_QUERY_TIMEOUT}
+            placeholder={`default = ${DEFAULT_QUERY_TIMEOUT}, 0 = unlimited`}
+            min={0}
+            displayValue={displayValueWithUnit('seconds')}
           />
           <ClickhouseSettingForm
             settingKey="metadataMaxRowsToRead"
@@ -1282,15 +1301,9 @@ function TeamQueryConfigSection() {
             tooltip="The maximum number of rows that can be read from a table when running a query"
             type="number"
             defaultValue={DEFAULT_METADATA_MAX_ROWS_TO_READ}
-            placeholder={`Enter value (default: ${DEFAULT_METADATA_MAX_ROWS_TO_READ.toLocaleString()}, 0 = unlimited)`}
+            placeholder={`default = ${DEFAULT_METADATA_MAX_ROWS_TO_READ.toLocaleString()}, 0 = unlimited`}
             min={0}
-            displayValue={value =>
-              value == null
-                ? `System Default (${DEFAULT_METADATA_MAX_ROWS_TO_READ.toLocaleString()})`
-                : value === 0
-                  ? 'Unlimited'
-                  : value.toLocaleString()
-            }
+            displayValue={displayValueWithUnit('rows')}
           />
           <ClickhouseSettingForm
             settingKey="fieldMetadataDisabled"

packages/app/src/clickhouse.ts

@@ -7,6 +7,7 @@
 
 import {
   chSql,
+  ClickhouseClientOptions,
   ColumnMeta,
   ResponseJSON,
 } from '@hyperdx/common-utils/dist/clickhouse';
@@ -16,28 +17,50 @@ import { useQuery, UseQueryOptions } from '@tanstack/react-query';
 import { IS_LOCAL_MODE } from '@/config';
 import { getLocalConnections } from '@/connection';
 
+import api from './api';
+import { DEFAULT_QUERY_TIMEOUT } from './defaults';
+
 const PROXY_CLICKHOUSE_HOST = '/api/clickhouse-proxy';
 
-export const getClickhouseClient = () => {
+export const getClickhouseClient = (
+  options: ClickhouseClientOptions = {},
+): ClickhouseClient => {
   if (IS_LOCAL_MODE) {
     const localConnections = getLocalConnections();
     if (localConnections.length === 0) {
       console.warn('No local connection found');
       return new ClickhouseClient({
         host: '',
+        ...options,
       });
     }
     return new ClickhouseClient({
       host: localConnections[0].host,
       username: localConnections[0].username,
       password: localConnections[0].password,
+      ...options,
     });
   }
   return new ClickhouseClient({
     host: PROXY_CLICKHOUSE_HOST,
+    ...options,
   });
 };
 
+export const useClickhouseClient = (
+  options: ClickhouseClientOptions = {},
+): ClickhouseClient => {
+  const { data: me } = api.useMe();
+  const teamQueryTimeout = me?.team?.queryTimeout;
+  if (teamQueryTimeout !== undefined) {
+    options.queryTimeout = teamQueryTimeout;
+  } else {
+    options.queryTimeout = DEFAULT_QUERY_TIMEOUT;
+  }
+
+  return getClickhouseClient(options);
+};
+
 export function useDatabasesDirect(
   { connectionId }: { connectionId: string },
   options?: Omit<UseQueryOptions<any, Error>, 'queryKey'>,

packages/app/src/components/KubeComponents.tsx

@@ -13,7 +13,7 @@ import {
 import { Anchor, Badge, Group, Text, Timeline } from '@mantine/core';
 import { useQuery, UseQueryOptions } from '@tanstack/react-query';
 
-import { getClickhouseClient } from '@/clickhouse';
+import { useClickhouseClient } from '@/clickhouse';
 import { getMetadata } from '@/metadata';
 import { getDisplayedTimestampValueExpression, getEventBody } from '@/source';
 
@@ -56,7 +56,7 @@ export const useV2LogBatch = <T = any,>(
   },
   options?: Omit<UseQueryOptions<any>, 'queryKey' | 'queryFn'>,
 ) => {
-  const clickhouseClient = getClickhouseClient();
+  const clickhouseClient = useClickhouseClient();
   return useQuery<ResponseJSON<T>, Error>({
     queryKey: [
       'v2LogBatch',

packages/app/src/defaults.ts

@@ -2,6 +2,7 @@ import type { ChartConfigWithDateRange } from '@hyperdx/common-utils/dist/types'
 
 // Limit defaults
 export const DEFAULT_SEARCH_ROW_LIMIT = 200;
+export const DEFAULT_QUERY_TIMEOUT = 60; // max_execution_time, seconds
 
 export function searchChartConfigDefaults(
   team: any | undefined | null,

packages/app/src/hooks/useChartConfig.tsx

@@ -14,7 +14,7 @@ import { format } from '@hyperdx/common-utils/dist/sqlFormatter';
 import { ChartConfigWithOptDateRange } from '@hyperdx/common-utils/dist/types';
 import { useQuery, UseQueryOptions } from '@tanstack/react-query';
 
-import { getClickhouseClient } from '@/clickhouse';
+import { useClickhouseClient } from '@/clickhouse';
 import { IS_MTVIEWS_ENABLED } from '@/config';
 import { buildMTViewSelectQuery } from '@/hdxMTViews';
 import { getMetadata } from '@/metadata';
@@ -29,7 +29,7 @@ export function useQueriedChartConfig(
   options?: Partial<UseQueryOptions<ResponseJSON<any>>> &
     AdditionalUseQueriedChartConfigOptions,
 ) {
-  const clickhouseClient = getClickhouseClient();
+  const clickhouseClient = useClickhouseClient();
   const query = useQuery<ResponseJSON<any>, ClickHouseQueryError | Error>({
     queryKey: [config],
     queryFn: async ({ signal }) => {

packages/app/src/hooks/useExplainQuery.tsx

@@ -2,7 +2,7 @@ import { renderChartConfig } from '@hyperdx/common-utils/dist/renderChartConfig'
 import { ChartConfigWithDateRange } from '@hyperdx/common-utils/dist/types';
 import { useQuery, UseQueryOptions } from '@tanstack/react-query';
 
-import { getClickhouseClient } from '@/clickhouse';
+import { useClickhouseClient } from '@/clickhouse';
 import { getMetadata } from '@/metadata';
 
 export function useExplainQuery(
@@ -13,7 +13,7 @@ export function useExplainQuery(
     ..._config,
     with: undefined,
   };
-  const clickhouseClient = getClickhouseClient();
+  const clickhouseClient = useClickhouseClient();
   const { data, isLoading, error } = useQuery({
     queryKey: ['explain', config],
     queryFn: async ({ signal }) => {

packages/app/src/hooks/useOffsetPaginatedQuery.tsx

@@ -15,12 +15,22 @@ import {
   useQueryClient,
 } from '@tanstack/react-query';
 
+import api from '@/api';
 import { getClickhouseClient } from '@/clickhouse';
 import { getMetadata } from '@/metadata';
 import { omit } from '@/utils';
 
-function queryKeyFn(prefix: string, config: ChartConfigWithDateRange) {
-  return [prefix, config] as const;
+type TQueryKey = readonly [
+  string,
+  ChartConfigWithDateRange,
+  number | undefined,
+];
+function queryKeyFn(
+  prefix: string,
+  config: ChartConfigWithDateRange,
+  queryTimeout?: number,
+): TQueryKey {
+  return [prefix, config, queryTimeout];
 }
 
 type TPageParam = number;
@@ -34,11 +44,12 @@ type TData = {
   pageParams: TPageParam[];
 };
 
-const queryFn: QueryFunction<
-  TQueryFnData,
-  readonly [string, ChartConfigWithDateRange],
-  number
-> = async ({ queryKey, pageParam, signal, meta }) => {
+const queryFn: QueryFunction<TQueryFnData, TQueryKey, number> = async ({
+  queryKey,
+  pageParam,
+  signal,
+  meta,
+}) => {
   if (meta == null) {
     throw new Error('Query missing client meta');
   }
@@ -60,7 +71,8 @@ const queryFn: QueryFunction<
     getMetadata(),
   );
 
-  const clickhouseClient = getClickhouseClient();
+  const queryTimeout = queryKey[2];
+  const clickhouseClient = getClickhouseClient({ queryTimeout });
   const resultSet =
     await clickhouseClient.query<'JSONCompactEachRowWithNamesAndTypes'>({
       query: query.sql,
@@ -247,7 +259,8 @@ export default function useOffsetPaginatedQuery(
     queryKeyPrefix?: string;
   } = {},
 ) {
-  const key = queryKeyFn(queryKeyPrefix, config);
+  const { data: meData } = api.useMe();
+  const key = queryKeyFn(queryKeyPrefix, config, meData?.team?.queryTimeout);
   const queryClient = useQueryClient();
   const matchedQueries = queryClient.getQueriesData<TData>({
     queryKey: [queryKeyPrefix, omit(config, ['dateRange'])],
@@ -268,7 +281,7 @@ export default function useOffsetPaginatedQuery(
     TQueryFnData,
     Error | ClickHouseQueryError,
     TData,
-    Readonly<[string, typeof config]>,
+    TQueryKey,
     TPageParam
   >({
     queryKey: key,

packages/app/src/metadata.ts

@@ -2,5 +2,8 @@ import { getMetadata as _getMetadata } from '@hyperdx/common-utils/dist/metadata
 
 import { getClickhouseClient } from '@/clickhouse';
 
+import { DEFAULT_QUERY_TIMEOUT } from './defaults';
+
 // TODO: Get rid of this function and convert to singleton
-export const getMetadata = () => _getMetadata(getClickhouseClient());
+export const getMetadata = () =>
+  _getMetadata(getClickhouseClient({ queryTimeout: DEFAULT_QUERY_TIMEOUT }));

packages/app/src/sessions.ts

@@ -14,7 +14,7 @@ import { useQuery, UseQueryOptions } from '@tanstack/react-query';
 import { getMetadata } from '@/metadata';
 import { usePrevious } from '@/utils';
 
-import { getClickhouseClient } from './clickhouse';
+import { getClickhouseClient, useClickhouseClient } from './clickhouse';
 import { IS_LOCAL_MODE } from './config';
 import { getLocalConnections } from './connection';
 import { useSource } from './source';
@@ -53,7 +53,7 @@ export function useSessions(
 ) {
   const FIXED_SDK_ATTRIBUTES = ['teamId', 'teamName', 'userEmail', 'userName'];
   const SESSIONS_CTE_NAME = 'sessions';
-  const clickhouseClient = getClickhouseClient();
+  const clickhouseClient = useClickhouseClient();
   return useQuery<ResponseJSON<Session>, Error>({
     queryKey: [
       'sessions',

packages/common-utils/src/clickhouse/browser.ts

@@ -81,7 +81,7 @@ export class ClickhouseClient extends BaseClickhouseClient {
     } catch (e) {
       debugSql = query;
     }
-    let _url = this.host;
+    let _url = this.host!;
 
     // eslint-disable-next-line no-console
     console.log('--------------------------------------------------------');
@@ -96,6 +96,12 @@ export class ClickhouseClient extends BaseClickhouseClient {
     if (clickhouse_settings?.max_rows_to_read && this.maxRowReadOnly) {
       delete clickhouse_settings['max_rows_to_read'];
     }
+    if (
+      clickhouse_settings?.max_execution_time === undefined &&
+      (this.queryTimeout || 0) > 0
+    ) {
+      clickhouse_settings.max_execution_time = this.queryTimeout;
+    }
 
     clickhouse_settings = {
       date_time_output_format: 'iso',

packages/common-utils/src/clickhouse/index.ts

@@ -367,15 +367,18 @@ export interface QueryInputs<Format extends DataFormat> {
 }
 
 export type ClickhouseClientOptions = {
-  host: string;
+  host?: string;
   username?: string;
   password?: string;
+  queryTimeout?: number;
 };
 
 export abstract class BaseClickhouseClient {
-  protected readonly host: string;
+  protected readonly host?: string;
   protected readonly username?: string;
   protected readonly password?: string;
+  protected readonly queryTimeout?: number;
+
   /*
    * Some clickhouse db's (the demo instance for example) make the
    * max_rows_to_read setting readonly and the query will fail if you try to
@@ -384,10 +387,16 @@ export abstract class BaseClickhouseClient {
   protected maxRowReadOnly: boolean;
   protected requestTimeout: number = 3600000; // TODO: make configurable
 
-  constructor({ host, username, password }: ClickhouseClientOptions) {
-    this.host = host;
+  constructor({
+    host,
+    username,
+    password,
+    queryTimeout,
+  }: ClickhouseClientOptions) {
+    this.host = host!;
     this.username = username;
     this.password = password;
+    this.queryTimeout = queryTimeout;
     this.maxRowReadOnly = false;
   }
 

packages/common-utils/src/clickhouse/node.ts

@@ -48,6 +48,12 @@ export class ClickhouseClient extends BaseClickhouseClient {
     if (clickhouse_settings?.max_rows_to_read && this.maxRowReadOnly) {
       delete clickhouse_settings['max_rows_to_read'];
     }
+    if (
+      clickhouse_settings?.max_execution_time === undefined &&
+      (this.queryTimeout || 0) > 0
+    ) {
+      clickhouse_settings.max_execution_time = this.queryTimeout;
+    }
 
     const _client = createClient({
       url: this.host,