diff --git a/.changeset/shaggy-clocks-whisper.md b/.changeset/shaggy-clocks-whisper.md new file mode 100644 index 000000000..f2d7e2b75 --- /dev/null +++ b/.changeset/shaggy-clocks-whisper.md @@ -0,0 +1,5 @@ +--- +"@hyperdx/api": minor +--- + +Allows defining the ClickHouse request timeout value from the command line on the check-alert task diff --git a/packages/api/src/tasks/__tests__/types.test.ts b/packages/api/src/tasks/__tests__/types.test.ts index d0ebb7648..e0224cce9 100644 --- a/packages/api/src/tasks/__tests__/types.test.ts +++ b/packages/api/src/tasks/__tests__/types.test.ts @@ -1,195 +1,94 @@ import { asTaskArgs } from '../types'; describe('asTaskArgs', () => { - it('should return valid TaskArgs for valid input', () => { - const validArgs = { - _: ['command', 'arg1', 'arg2'], - provider: 'default', - }; - - const result = asTaskArgs(validArgs); - - expect(result).toEqual({ - taskName: 'command', - provider: 'default', + describe('invalid inputs', () => { + it('should throw error for null input', () => { + expect(() => asTaskArgs(null)).toThrow( + 'Arguments cannot be null or undefined', + ); }); - expect(result.taskName).toBe('command'); - // For non-check-alerts tasks, we need to use type assertion to access provider - expect((result as any).provider).toBe('default'); - }); - - it('should return valid TaskArgs when provider is undefined', () => { - const validArgs = { - _: ['command'], - }; - const result = asTaskArgs(validArgs); - - expect(result).toEqual({ - taskName: 'command', - provider: undefined, + it('should throw error for undefined input', () => { + expect(() => asTaskArgs(undefined)).toThrow( + 'Arguments cannot be null or undefined', + ); }); - expect(result.taskName).toBe('command'); - // For non-check-alerts tasks, we need to use type assertion to access provider - expect((result as any).provider).toBeUndefined(); - }); - - it('should throw error for null input', () => { - expect(() => asTaskArgs(null)).toThrow( - 'Arguments cannot be null or undefined', - ); - }); - - it('should throw error for undefined input', () => { - expect(() => asTaskArgs(undefined)).toThrow( - 'Arguments cannot be null or undefined', - ); - }); - - it('should throw error for non-object input', () => { - expect(() => asTaskArgs('string')).toThrow('Arguments must be an object'); - expect(() => asTaskArgs(123)).toThrow('Arguments must be an object'); - expect(() => asTaskArgs(true)).toThrow('Arguments must be an object'); - expect(() => asTaskArgs(false)).toThrow('Arguments must be an object'); - expect(() => asTaskArgs([])).toThrow('Arguments must be an object'); - }); - - it('should throw error when _ property is missing', () => { - const invalidArgs = { - provider: 'default', - }; - - expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Arguments must have a "_" property that is an array', - ); - }); - - it('should throw error when _ property is not an array', () => { - const invalidArgs = { - _: 'not an array', - provider: 'default', - }; - - expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Arguments must have a "_" property that is an array', - ); - }); - it('should throw error when _ property is null', () => { - const invalidArgs = { - _: null, - provider: 'default', - }; - - expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Arguments must have a "_" property that is an array', - ); - }); - - it('should handle empty array for _ property', () => { - const validArgs = { - _: [], - provider: 'default', - }; - - const result = asTaskArgs(validArgs); - - expect(result).toEqual({ - taskName: undefined, - provider: 'default', + it('should throw error for non-object input', () => { + expect(() => asTaskArgs('string')).toThrow('Arguments must be an object'); + expect(() => asTaskArgs(123)).toThrow('Arguments must be an object'); + expect(() => asTaskArgs(true)).toThrow('Arguments must be an object'); + expect(() => asTaskArgs(false)).toThrow('Arguments must be an object'); + expect(() => asTaskArgs([])).toThrow('Arguments must be an object'); }); - expect(result.taskName).toBeUndefined(); - }); - - it('should accept array with only strings for _ property', () => { - const validArgs = { - _: ['string', '123', 'true', 'null', 'undefined'], - provider: 'default', - }; - const result = asTaskArgs(validArgs); + it('should throw error when _ property is missing', () => { + const invalidArgs = { + provider: 'default', + }; - expect(result).toEqual({ - taskName: 'string', - provider: 'default', + expect(() => asTaskArgs(invalidArgs)).toThrow( + 'Arguments must have a "_" property that is an array', + ); }); - expect(result.taskName).toBe('string'); - }); - it('should extract taskName from first argument', () => { - const validArgs = { - _: ['command'], - provider: 'default', - extraProperty: 'value', - anotherProperty: 123, - }; - - const result = asTaskArgs(validArgs); + it('should throw error when _ property is not an array', () => { + const invalidArgs = { + _: 'not an array', + provider: 'default', + }; - expect(result).toEqual({ - taskName: 'command', - provider: 'default', + expect(() => asTaskArgs(invalidArgs)).toThrow( + 'Arguments must have a "_" property that is an array', + ); }); - expect(result.taskName).toBe('command'); - }); - it('should accept check-alerts task with provider', () => { - const validArgs = { - _: ['check-alerts'], - provider: 'default', - }; - - const result = asTaskArgs(validArgs); + it('should throw error when _ property is null', () => { + const invalidArgs = { + _: null, + provider: 'default', + }; - expect(result).toEqual({ - taskName: 'check-alerts', - provider: 'default', + expect(() => asTaskArgs(invalidArgs)).toThrow( + 'Arguments must have a "_" property that is an array', + ); }); - expect(result.taskName).toBe('check-alerts'); - // For check-alerts tasks, provider property is directly accessible - if (result.taskName === 'check-alerts') { - expect(result.provider).toBe('default'); - } - }); - - it('should accept check-alerts task without provider', () => { - const validArgs = { - _: ['check-alerts'], - }; - const result = asTaskArgs(validArgs); + it('should throw error when _ is empty', () => { + const validArgs = { + _: [], + provider: 'default', + }; - expect(result).toEqual({ - taskName: 'check-alerts', - provider: undefined, + expect(() => asTaskArgs(validArgs)).toThrow( + 'Task name needs to be specified', + ); }); - expect(result.taskName).toBe('check-alerts'); - // For check-alerts tasks, provider property is directly accessible - if (result.taskName === 'check-alerts') { - expect(result.provider).toBeUndefined(); - } }); - it('should accept ping-pong task without provider', () => { - const validArgs = { - _: ['ping-pong'], - }; + describe('check-alerts task', () => { + it('should accept check-alerts task with provider', () => { + const validArgs = { + _: ['check-alerts'], + provider: 'default', + }; - const result = asTaskArgs(validArgs); + const result = asTaskArgs(validArgs); - expect(result).toEqual({ - taskName: 'ping-pong', + expect(result).toEqual({ + taskName: 'check-alerts', + provider: 'default', + }); + expect(result.taskName).toBe('check-alerts'); + // For check-alerts tasks, provider property is directly accessible + if (result.taskName === 'check-alerts') { + expect(result.provider).toBe('default'); + } }); - expect(result.taskName).toBe('ping-pong'); - // Ping-pong tasks should not have a provider property - expect('provider' in result).toBe(false); - }); - describe('concurrency parameter validation', () => { - it('should accept check-alerts task with valid concurrency', () => { + it('should accept check-alerts task without provider', () => { const validArgs = { _: ['check-alerts'], - concurrency: 4, }; const result = asTaskArgs(validArgs); @@ -197,18 +96,18 @@ describe('asTaskArgs', () => { expect(result).toEqual({ taskName: 'check-alerts', provider: undefined, - concurrency: 4, }); expect(result.taskName).toBe('check-alerts'); + // For check-alerts tasks, provider property is directly accessible if (result.taskName === 'check-alerts') { - expect(result.concurrency).toBe(4); + expect(result.provider).toBeUndefined(); } }); - it('should accept check-alerts task with concurrency value of 1', () => { + it('should accept check-alerts task with valid concurrency', () => { const validArgs = { _: ['check-alerts'], - concurrency: 1, + concurrency: 4, }; const result = asTaskArgs(validArgs); @@ -216,18 +115,18 @@ describe('asTaskArgs', () => { expect(result).toEqual({ taskName: 'check-alerts', provider: undefined, - concurrency: 1, + concurrency: 4, }); expect(result.taskName).toBe('check-alerts'); if (result.taskName === 'check-alerts') { - expect(result.concurrency).toBe(1); + expect(result.concurrency).toBe(4); } }); - it('should accept check-alerts task with large concurrency values', () => { + it('should accept check-alerts task with concurrency value of 1', () => { const validArgs = { _: ['check-alerts'], - concurrency: 100, + concurrency: 1, }; const result = asTaskArgs(validArgs); @@ -235,30 +134,30 @@ describe('asTaskArgs', () => { expect(result).toEqual({ taskName: 'check-alerts', provider: undefined, - concurrency: 100, + concurrency: 1, }); expect(result.taskName).toBe('check-alerts'); if (result.taskName === 'check-alerts') { - expect(result.concurrency).toBe(100); + expect(result.concurrency).toBe(1); } }); - it('should accept check-alerts task without concurrency parameter', () => { + it('should accept check-alerts task with a concurrency value', () => { const validArgs = { _: ['check-alerts'], - provider: 'default', + concurrency: 100, }; const result = asTaskArgs(validArgs); expect(result).toEqual({ taskName: 'check-alerts', - provider: 'default', - concurrency: undefined, + provider: undefined, + concurrency: 100, }); expect(result.taskName).toBe('check-alerts'); if (result.taskName === 'check-alerts') { - expect(result.concurrency).toBeUndefined(); + expect(result.concurrency).toBe(100); } }); @@ -268,125 +167,118 @@ describe('asTaskArgs', () => { concurrency: 'invalid', }; - expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Concurrency must be a number if provided', - ); + expect(() => asTaskArgs(invalidArgs)).toThrow('Expected number'); }); - it('should throw error when concurrency is boolean', () => { + it('should throw error when concurrency is zero', () => { const invalidArgs = { _: ['check-alerts'], - concurrency: true, + concurrency: 0, }; expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Concurrency must be a number if provided', + 'concurrency must be at least 1', ); }); - it('should throw error when concurrency is null', () => { + it('should throw error when concurrency is negative', () => { const invalidArgs = { _: ['check-alerts'], - concurrency: null, + concurrency: -1, }; expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Concurrency must be a number if provided', + 'concurrency must be at least 1', ); }); - it('should throw error when concurrency is an object', () => { + it('should throw error when concurrency is a positive decimal', () => { const invalidArgs = { _: ['check-alerts'], - concurrency: { value: 4 }, + concurrency: 2.5, }; expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Concurrency must be a number if provided', + 'concurrency must be an integer', ); }); - it('should throw error when concurrency is an array', () => { - const invalidArgs = { + it('should accept check-alerts task with valid sourceTimeoutMs', () => { + const validArgs = { _: ['check-alerts'], - concurrency: [4], + sourceTimeoutMs: 5000, }; - expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Concurrency must be a number if provided', - ); - }); - - it('should throw error when concurrency is zero', () => { - const invalidArgs = { - _: ['check-alerts'], - concurrency: 0, - }; + const result = asTaskArgs(validArgs); - expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Concurrency cannot be less than 1', - ); + expect(result).toEqual({ + taskName: 'check-alerts', + provider: undefined, + concurrency: undefined, + sourceTimeoutMs: 5000, + }); + expect(result.taskName).toBe('check-alerts'); + if (result.taskName === 'check-alerts') { + expect(result.sourceTimeoutMs).toBe(5000); + } }); - it('should throw error when concurrency is negative', () => { - const invalidArgs = { + it('should accept check-alerts task with sourceTimeoutMs of 0', () => { + const validArgs = { _: ['check-alerts'], - concurrency: -1, + sourceTimeoutMs: 0, }; - expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Concurrency cannot be less than 1', - ); - }); - - it('should throw error when concurrency is a negative decimal', () => { - const invalidArgs = { - _: ['check-alerts'], - concurrency: -0.5, - }; + const result = asTaskArgs(validArgs); - expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Concurrency must be an integer if provided', - ); + expect(result).toEqual({ + taskName: 'check-alerts', + provider: undefined, + concurrency: undefined, + sourceTimeoutMs: 0, + }); + expect(result.taskName).toBe('check-alerts'); + if (result.taskName === 'check-alerts') { + expect(result.sourceTimeoutMs).toBe(0); + } }); - it('should throw error when concurrency is a positive decimal', () => { + it('should throw error when sourceTimeoutMs is not a number', () => { const invalidArgs = { _: ['check-alerts'], - concurrency: 2.5, + sourceTimeoutMs: 'invalid', }; - expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Concurrency must be an integer if provided', - ); + expect(() => asTaskArgs(invalidArgs)).toThrow('Expected number'); }); - it('should throw error when concurrency is a small decimal', () => { + it('should throw error when sourceTimeoutMs is negative', () => { const invalidArgs = { _: ['check-alerts'], - concurrency: 1.1, + sourceTimeoutMs: -1, }; expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Concurrency must be an integer if provided', + 'sourceTimeoutMs must be a non-negative value', ); }); - it('should throw error when concurrency is a large decimal', () => { + it('should throw error when sourceTimeoutMs is a small decimal', () => { const invalidArgs = { _: ['check-alerts'], - concurrency: 100.999, + sourceTimeoutMs: 1.1, }; expect(() => asTaskArgs(invalidArgs)).toThrow( - 'Concurrency must be an integer if provided', + 'sourceTimeoutMs must be an int', ); }); + }); - it('should ignore concurrency parameter for non-check-alerts tasks', () => { + describe('ping-pong task', () => { + it('should accept ping-pong task without provider', () => { const validArgs = { _: ['ping-pong'], - concurrency: 4, }; const result = asTaskArgs(validArgs); @@ -395,25 +287,8 @@ describe('asTaskArgs', () => { taskName: 'ping-pong', }); expect(result.taskName).toBe('ping-pong'); - // Ping-pong tasks should not have a concurrency property - expect('concurrency' in result).toBe(false); - }); - - it('should ignore concurrency parameter for unknown task types', () => { - const validArgs = { - _: ['unknown-task'], - concurrency: 4, - }; - - const result = asTaskArgs(validArgs); - - expect(result).toEqual({ - taskName: 'unknown-task', - provider: undefined, - }); - expect(result.taskName).toBe('unknown-task'); - // Unknown task types should not process concurrency parameter - expect('concurrency' in result).toBe(false); + // Ping-pong tasks should not have a provider property + expect('provider' in result).toBe(false); }); }); }); diff --git a/packages/api/src/tasks/checkAlerts.ts b/packages/api/src/tasks/checkAlerts.ts index 41860be0d..1644ee32d 100644 --- a/packages/api/src/tasks/checkAlerts.ts +++ b/packages/api/src/tasks/checkAlerts.ts @@ -448,7 +448,10 @@ export default class CheckAlertTask implements HdxTask { alertCount: alerts.length, }); - const clickhouseClient = await this.provider.getClickHouseClient(conn); + const clickhouseClient = await this.provider.getClickHouseClient( + conn, + this.args.sourceTimeoutMs, + ); for (const alert of alerts) { await this.task_queue.add(() => diff --git a/packages/api/src/tasks/providers/default.ts b/packages/api/src/tasks/providers/default.ts index f2f2e25c3..0d6833b1b 100644 --- a/packages/api/src/tasks/providers/default.ts +++ b/packages/api/src/tasks/providers/default.ts @@ -297,12 +297,10 @@ export default class DefaultAlertProvider implements AlertProvider { return new Map(webhooks.map(w => [w.id, w])); } - async getClickHouseClient({ - host, - username, - password, - id, - }: IConnection): Promise { + async getClickHouseClient( + { host, username, password, id }: IConnection, + requestTimeout?: number, + ): Promise { if (!password && password !== '') { logger.info({ message: `connection password not found`, @@ -316,6 +314,7 @@ export default class DefaultAlertProvider implements AlertProvider { username, password, application: `hyperdx-alerts ${config.CODE_VERSION}`, + requestTimeout: requestTimeout ?? 30_000, }); } } diff --git a/packages/api/src/tasks/providers/index.ts b/packages/api/src/tasks/providers/index.ts index b1e8a8d06..a13416159 100644 --- a/packages/api/src/tasks/providers/index.ts +++ b/packages/api/src/tasks/providers/index.ts @@ -14,6 +14,7 @@ import DefaultAlertProvider from '@/tasks/providers/default'; import logger from '@/utils/logger'; import { AggregatedAlertHistory } from '../checkAlerts'; +import { CheckAlertsTaskArgs } from '../types'; export enum AlertTaskType { SAVED_SEARCH, @@ -81,7 +82,10 @@ export interface AlertProvider { getWebhooks(teamId: string | ObjectId): Promise>; /** Create and return an authenticated ClickHouse client */ - getClickHouseClient(connection: IConnection): Promise; + getClickHouseClient( + connection: IConnection, + requestTimeout?: number, + ): Promise; } const ADDITIONAL_PROVIDERS: Record AlertProvider> = { diff --git a/packages/api/src/tasks/types.ts b/packages/api/src/tasks/types.ts index 7445c7723..d5d705cfa 100644 --- a/packages/api/src/tasks/types.ts +++ b/packages/api/src/tasks/types.ts @@ -1,18 +1,37 @@ +import { z } from 'zod'; + /** * Command line arguments structure for tasks. * Contains task name and optional provider configuration. */ -export type PingTaskArgs = { taskName: 'ping-pong' }; -export type CheckAlertsTaskArgs = { - taskName: 'check-alerts'; - // name of the provider module to use for fetching alert task data. If not defined, - // the default provider will be used. - provider?: string; - // Limits number of concurrent tasks processed. If omitted, there is no concurrency - // limit. Must be an integer greater than 0. - concurrency?: number; -}; -export type TaskArgs = PingTaskArgs | CheckAlertsTaskArgs; +const pingTaskArgsSchema = z.object({ + taskName: z.literal('ping-pong'), +}); + +const checkAlertsTaskArgsSchema = z.object({ + taskName: z.literal('check-alerts'), + provider: z.string().optional(), + concurrency: z + .number() + .int('concurrency must be an integer') + .min(1, 'concurrency must be at least 1') + .max(1024, 'concurrency must be less than 1024') + .optional(), + sourceTimeoutMs: z + .number() + .int('sourceTimeoutMs must be an int') + .nonnegative('sourceTimeoutMs must be a non-negative value') + .optional(), +}); + +const taskArgsSchema = z.discriminatedUnion('taskName', [ + pingTaskArgsSchema, + checkAlertsTaskArgsSchema, +]); + +export type PingTaskArgs = z.infer; +export type CheckAlertsTaskArgs = z.infer; +export type TaskArgs = z.infer; /** * Validates and converts command line arguments to TaskArgs type. @@ -41,49 +60,16 @@ export function asTaskArgs(argv: any): TaskArgs { throw new Error('All arguments in "_" array must be strings'); } - const taskName = argv._[0]; - if (taskName === 'check-alerts') { - const { provider, concurrency } = argv; - if (provider) { - if (typeof provider !== 'string') { - throw new Error('Provider must be a string if provided'); - } - - if (provider.trim() === '') { - throw new Error('Provider must contain valid characters'); - } - } - - if (concurrency !== undefined) { - if (typeof concurrency !== 'number') { - throw new Error('Concurrency must be a number if provided'); - } - - if (!Number.isInteger(concurrency)) { - throw new Error('Concurrency must be an integer if provided'); - } - - if (concurrency < 1) { - throw new Error('Concurrency cannot be less than 1'); - } - } - - return { - taskName: 'check-alerts', - provider: provider, - concurrency: concurrency, - }; - } else if (taskName === 'ping-pong') { - return { - taskName: 'ping-pong', - }; - } else { - // For any other task names, create a generic structure without provider - return { - taskName, - provider: argv.provider, - } as TaskArgs; + const { _, ...rest } = argv; + if (_.length < 1) { + throw new Error('Task name needs to be specified'); } + const taskName = _[0]; + + return taskArgsSchema.parse({ + taskName, + ...rest, + }); } /** diff --git a/packages/common-utils/src/clickhouse/index.ts b/packages/common-utils/src/clickhouse/index.ts index cc87a621e..aca6850e2 100644 --- a/packages/common-utils/src/clickhouse/index.ts +++ b/packages/common-utils/src/clickhouse/index.ts @@ -408,6 +408,8 @@ export type ClickhouseClientOptions = { queryTimeout?: number; /** Application name, used as the client's HTTP user-agent header */ application?: string; + /** Defines how long the client will wait for a response from the ClickHouse server before aborting the request, in milliseconds */ + requestTimeout?: number; }; export abstract class BaseClickhouseClient { @@ -423,7 +425,7 @@ export abstract class BaseClickhouseClient { * query with max_rows_to_read specified */ protected maxRowReadOnly: boolean; - protected requestTimeout: number = 3600000; // TODO: make configurable + protected requestTimeout: number = 3600000; constructor({ host, @@ -431,6 +433,7 @@ export abstract class BaseClickhouseClient { password, queryTimeout, application, + requestTimeout, }: ClickhouseClientOptions) { this.host = host!; this.username = username; @@ -438,6 +441,9 @@ export abstract class BaseClickhouseClient { this.queryTimeout = queryTimeout; this.maxRowReadOnly = false; this.application = application; + if (requestTimeout != null && requestTimeout >= 0) { + this.requestTimeout = requestTimeout; + } } protected getClient(): WebClickHouseClient | NodeClickHouseClient {