Allow to specify additional HTTP headers for ClickHouse requests (#224)

packages/client-common/src/client.ts

@@ -92,6 +92,7 @@ export interface ClickHouseClientConfigOptions<Stream> {
     level?: ClickHouseLogLevel
   }
   session_id?: string
+  additional_headers?: Record<string, number | string | string[]>
 }
 
 export type BaseClickHouseClientConfigOptions<Stream> = Omit<
@@ -335,6 +336,7 @@ function getConnectionParams<Stream>(
         : new DefaultLogger(),
       config.log?.level
     ),
+    additional_headers: config.additional_headers,
   }
 }
 

packages/client-common/src/connection.ts

@@ -16,6 +16,7 @@ export interface ConnectionParams {
   clickhouse_settings: ClickHouseSettings
   logWriter: LogWriter
   application_id?: string
+  additional_headers?: Record<string, number | string | string[]>
 }
 
 export interface ConnBaseQueryParams {

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<ConnectionParams>) {
+    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,
+    })
+  }
+})

packages/client-node/__tests__/unit/node_http_adapter.test.ts

@@ -563,7 +563,7 @@ describe('[Node.js] HttpAdapter', () => {
   function buildHttpAdapter(config: Partial<ConnectionParams>) {
     return new NodeHttpConnection({
       ...{
-        url: new URL('http://localhost:8132'),
+        url: new URL('http://localhost:8123'),
 
         connect_timeout: 10_000,
         request_timeout: 30_000,

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<string, number | string | string[]>
   ): Http.OutgoingHttpHeaders {
     return {
       Authorization: `Basic ${Buffer.from(`${username}:${password}`).toString(
         'base64'
       )}`,
       'User-Agent': getUserAgent(this.params.application_id),
+      ...additional_headers,
     }
   }
 

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<string, number | string | string[]>
   ): 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 {

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<typeof window.fetch>
+    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<typeof window.fetch>
     beforeEach(() => {

packages/client-web/src/connection/web_connection.ts

@@ -35,6 +35,7 @@ export class WebConnection implements Connection<ReadableStream> {
   constructor(private readonly params: WebConnectionParams) {
     this.defaultHeaders = {
       Authorization: `Basic ${btoa(`${params.username}:${params.password}`)}`,
+      ...params?.additional_headers,
     }
   }