From 9b358ea37ab25ef7dbfe226e9eaa4c573ff0b032 Mon Sep 17 00:00:00 2001 From: teawithfruit Date: Fri, 19 Jan 2024 17:47:41 +0100 Subject: [PATCH] Allow to specify additional HTTP headers for ClickHouse requests (#224) --- packages/client-common/src/client.ts | 2 + packages/client-common/src/connection.ts | 1 + .../integration/node_connection.test.ts | 102 ++++++++++++++++++ .../__tests__/unit/node_http_adapter.test.ts | 2 +- .../src/connection/node_base_connection.ts | 10 +- .../src/connection/node_https_connection.ts | 5 +- .../integration/web_connection.test.ts | 26 +++++ .../src/connection/web_connection.ts | 1 + 8 files changed, 144 insertions(+), 5 deletions(-) create mode 100644 packages/client-node/__tests__/integration/node_connection.test.ts diff --git a/packages/client-common/src/client.ts b/packages/client-common/src/client.ts index 7936f105..ab0e5e45 100644 --- a/packages/client-common/src/client.ts +++ b/packages/client-common/src/client.ts @@ -92,6 +92,7 @@ export interface ClickHouseClientConfigOptions { level?: ClickHouseLogLevel } session_id?: string + additional_headers?: Record } export type BaseClickHouseClientConfigOptions = Omit< @@ -335,6 +336,7 @@ function getConnectionParams( : new DefaultLogger(), config.log?.level ), + additional_headers: config.additional_headers, } } diff --git a/packages/client-common/src/connection.ts b/packages/client-common/src/connection.ts index a55885bd..d99df511 100644 --- a/packages/client-common/src/connection.ts +++ b/packages/client-common/src/connection.ts @@ -16,6 +16,7 @@ export interface ConnectionParams { clickhouse_settings: ClickHouseSettings logWriter: LogWriter application_id?: string + additional_headers?: Record } export interface ConnBaseQueryParams { diff --git a/packages/client-node/__tests__/integration/node_connection.test.ts b/packages/client-node/__tests__/integration/node_connection.test.ts new file mode 100644 index 00000000..385f568e --- /dev/null +++ b/packages/client-node/__tests__/integration/node_connection.test.ts @@ -0,0 +1,102 @@ +import type { ConnectionParams } from '@clickhouse/client-common' +import { LogWriter } from '@clickhouse/client-common' +import { TestLogger } from '@test/utils' +import { randomUUID } from '@test/utils/guid' +import type { ClientRequest } from 'http' +import Http from 'http' +import Stream from 'stream' +import { NodeHttpConnection } from '../../src/connection' + +describe('[Node.js] Connection', () => { + it('should be possible to set additional_headers', async () => { + const request = stubClientRequest() + const httpRequestStub = spyOn(Http, 'request').and.returnValue(request) + const adapter = buildHttpAdapter({ + additional_headers: { + 'Test-Header': 'default', + }, + }) + + const selectPromise = adapter.query({ + query: 'SELECT * FROM system.numbers LIMIT 5', + }) + + const responseBody = 'foobar' + request.emit( + 'response', + buildIncomingMessage({ + body: responseBody, + }) + ) + + await selectPromise + + expect(httpRequestStub).toHaveBeenCalledTimes(1) + const calledWith = httpRequestStub.calls.mostRecent().args[1] + expect(calledWith.headers?.['Test-Header']).toBe('default') + }) + + function buildIncomingMessage({ + body = '', + statusCode = 200, + headers = {}, + }: { + body?: string | Buffer + statusCode?: number + headers?: Http.IncomingHttpHeaders + }): Http.IncomingMessage { + const response = new Stream.Readable({ + read() { + this.push(body) + this.push(null) + }, + }) as Http.IncomingMessage + + response.statusCode = statusCode + response.headers = { + 'x-clickhouse-query-id': randomUUID(), + ...headers, + } + return response + } + + function stubClientRequest() { + const request = new Stream.Writable({ + write() { + /** stub */ + }, + }) as ClientRequest + request.getHeaders = () => ({}) + return request + } + + function buildHttpAdapter(config: Partial) { + return new NodeHttpConnection({ + ...{ + url: new URL('http://localhost:8123'), + + connect_timeout: 10_000, + request_timeout: 30_000, + compression: { + decompress_response: true, + compress_request: false, + }, + max_open_connections: Infinity, + + username: '', + password: '', + database: '', + clickhouse_settings: {}, + additional_headers: {}, + + logWriter: new LogWriter(new TestLogger()), + keep_alive: { + enabled: true, + socket_ttl: 2500, + retry_on_expired_socket: false, + }, + }, + ...config, + }) + } +}) diff --git a/packages/client-node/__tests__/unit/node_http_adapter.test.ts b/packages/client-node/__tests__/unit/node_http_adapter.test.ts index 3aa2a42e..28911c56 100644 --- a/packages/client-node/__tests__/unit/node_http_adapter.test.ts +++ b/packages/client-node/__tests__/unit/node_http_adapter.test.ts @@ -563,7 +563,7 @@ describe('[Node.js] HttpAdapter', () => { function buildHttpAdapter(config: Partial) { return new NodeHttpConnection({ ...{ - url: new URL('http://localhost:8132'), + url: new URL('http://localhost:8123'), connect_timeout: 10_000, request_timeout: 30_000, diff --git a/packages/client-node/src/connection/node_base_connection.ts b/packages/client-node/src/connection/node_base_connection.ts index d4a458e0..c24c51f8 100644 --- a/packages/client-node/src/connection/node_base_connection.ts +++ b/packages/client-node/src/connection/node_base_connection.ts @@ -82,18 +82,24 @@ export abstract class NodeBaseConnection this.logger = params.logWriter this.retry_expired_sockets = params.keep_alive.enabled && params.keep_alive.retry_on_expired_socket - this.headers = this.buildDefaultHeaders(params.username, params.password) + this.headers = this.buildDefaultHeaders( + params.username, + params.password, + params.additional_headers + ) } protected buildDefaultHeaders( username: string, - password: string + password: string, + additional_headers?: Record ): Http.OutgoingHttpHeaders { return { Authorization: `Basic ${Buffer.from(`${username}:${password}`).toString( 'base64' )}`, 'User-Agent': getUserAgent(this.params.application_id), + ...additional_headers, } } diff --git a/packages/client-node/src/connection/node_https_connection.ts b/packages/client-node/src/connection/node_https_connection.ts index 34320074..3b07c3a5 100644 --- a/packages/client-node/src/connection/node_https_connection.ts +++ b/packages/client-node/src/connection/node_https_connection.ts @@ -26,7 +26,8 @@ export class NodeHttpsConnection protected override buildDefaultHeaders( username: string, - password: string + password: string, + additional_headers?: Record ): Http.OutgoingHttpHeaders { if (this.params.tls?.type === 'Mutual') { return { @@ -41,7 +42,7 @@ export class NodeHttpsConnection 'X-ClickHouse-Key': password, } } - return super.buildDefaultHeaders(username, password) + return super.buildDefaultHeaders(username, password, additional_headers) } protected createClientRequest(params: RequestParams): Http.ClientRequest { diff --git a/packages/client-web/__tests__/integration/web_connection.test.ts b/packages/client-web/__tests__/integration/web_connection.test.ts index b71121bb..d01249fc 100644 --- a/packages/client-web/__tests__/integration/web_connection.test.ts +++ b/packages/client-web/__tests__/integration/web_connection.test.ts @@ -2,6 +2,32 @@ import { createClient } from '../../src' import type { WebClickHouseClient } from '../../src/client' describe('[Web] Connection', () => { + describe('additional_headers', () => { + let fetchSpy: jasmine.Spy + beforeEach(() => { + fetchSpy = spyOn(window, 'fetch').and.returnValue( + Promise.resolve(new Response()) + ) + }) + + it('should be possible to set', async () => { + const client = createClient({ + additional_headers: { + 'Test-Header': 'default', + }, + }) + const fetchParams = await pingAndGetRequestInit(client) + expect(fetchParams!.headers?.['Test-Header']).toBe('default') + + async function pingAndGetRequestInit(client: WebClickHouseClient) { + await client.ping() + expect(fetchSpy).toHaveBeenCalledTimes(1) + const [, fetchParams] = fetchSpy.calls.mostRecent().args + return fetchParams! + } + }) + }) + describe('KeepAlive setting', () => { let fetchSpy: jasmine.Spy beforeEach(() => { diff --git a/packages/client-web/src/connection/web_connection.ts b/packages/client-web/src/connection/web_connection.ts index 710736a7..a0fc1cbb 100644 --- a/packages/client-web/src/connection/web_connection.ts +++ b/packages/client-web/src/connection/web_connection.ts @@ -35,6 +35,7 @@ export class WebConnection implements Connection { constructor(private readonly params: WebConnectionParams) { this.defaultHeaders = { Authorization: `Basic ${btoa(`${params.username}:${params.password}`)}`, + ...params?.additional_headers, } }