Ability to configure the clickhouse keep alive settings

Author: ericallam

apps/webapp/app/env.server.ts

@@ -769,6 +769,8 @@ const EnvironmentSchema = z.object({
   RUN_REPLICATION_LEADER_LOCK_ADDITIONAL_TIME_MS: z.coerce.number().int().default(10_000),
   RUN_REPLICATION_LEADER_LOCK_RETRY_INTERVAL_MS: z.coerce.number().int().default(500),
   RUN_REPLICATION_WAIT_FOR_ASYNC_INSERT: z.string().default("0"),
+  RUN_REPLICATION_KEEP_ALIVE_ENABLED: z.string().default("1"),
+  RUN_REPLICATION_KEEP_ALIVE_IDLE_SOCKET_TTL_MS: z.coerce.number().int().default(9_000),
 });
 
 export type Environment = z.infer<typeof EnvironmentSchema>;

apps/webapp/app/services/runsReplicationInstance.server.ts

@@ -23,6 +23,11 @@ function initializeRunsReplicationInstance() {
   const clickhouse = new ClickHouse({
     url: env.RUN_REPLICATION_CLICKHOUSE_URL,
     name: "runs-replication",
+    keepAlive: {
+      enabled: env.RUN_REPLICATION_KEEP_ALIVE_ENABLED === "1",
+      idleSocketTtl: env.RUN_REPLICATION_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
+    },
+    logLevel: env.RUN_REPLICATION_LOG_LEVEL,
   });
 
   const service = new RunsReplicationService({

internal-packages/clickhouse/src/client/client.ts

@@ -15,14 +15,22 @@ import type {
   ClickhouseWriter,
 } from "./types.js";
 import { generateErrorMessage } from "zod-error";
-import { Logger } from "@trigger.dev/core/logger";
+import { Logger, type LogLevel } from "@trigger.dev/core/logger";
+import type { Agent as HttpAgent } from "http";
+import type { Agent as HttpsAgent } from "https";
 
 export type ClickhouseConfig = {
   name: string;
   url: string;
   tracer?: Tracer;
+  keepAlive?: {
+    enabled?: boolean;
+    idleSocketTtl?: number;
+  };
+  httpAgent?: HttpAgent | HttpsAgent;
   clickhouseSettings?: ClickHouseSettings;
   logger?: Logger;
+  logLevel?: LogLevel;
 };
 
 export class ClickhouseClient implements ClickhouseReader, ClickhouseWriter {
@@ -33,11 +41,12 @@ export class ClickhouseClient implements ClickhouseReader, ClickhouseWriter {
 
   constructor(config: ClickhouseConfig) {
     this.name = config.name;
-    this.logger = config.logger ?? new Logger("ClickhouseClient", "debug");
+    this.logger = config.logger ?? new Logger("ClickhouseClient", config.logLevel ?? "info");
 
     this.client = createClient({
       url: config.url,
-
+      keep_alive: config.keepAlive,
+      http_agent: config.httpAgent,
       clickhouse_settings: {
         ...config.clickhouseSettings,
         output_format_json_quote_64bit_integers: 0,

internal-packages/clickhouse/src/index.ts

@@ -3,29 +3,38 @@ import { ClickhouseClient } from "./client/client.js";
 import { ClickhouseReader, ClickhouseWriter } from "./client/types.js";
 import { NoopClient } from "./client/noop.js";
 import { insertTaskRuns, insertRawTaskRunPayloads } from "./taskRuns.js";
-import { Logger } from "@trigger.dev/core/logger";
+import { Logger, type LogLevel } from "@trigger.dev/core/logger";
+import type { Agent as HttpAgent } from "http";
+import type { Agent as HttpsAgent } from "https";
 
 export type * from "./taskRuns.js";
 
+export type ClickhouseCommonConfig = {
+  keepAlive?: {
+    enabled?: boolean;
+    idleSocketTtl?: number;
+  };
+  httpAgent?: HttpAgent | HttpsAgent;
+  clickhouseSettings?: ClickHouseSettings;
+  logger?: Logger;
+  logLevel?: LogLevel;
+};
+
 export type ClickHouseConfig =
-  | {
+  | ({
       name?: string;
       url?: string;
       writerUrl?: never;
       readerUrl?: never;
-      clickhouseSettings?: ClickHouseSettings;
-      logger?: Logger;
-    }
-  | {
+    } & ClickhouseCommonConfig)
+  | ({
       name?: never;
       url?: never;
       writerName?: string;
       writerUrl: string;
       readerName?: string;
       readerUrl: string;
-      clickhouseSettings?: ClickHouseSettings;
-      logger?: Logger;
-    };
+    } & ClickhouseCommonConfig);
 
 export class ClickHouse {
   public readonly reader: ClickhouseReader;
@@ -47,6 +56,9 @@ export class ClickHouse {
         url: config.url,
         clickhouseSettings: config.clickhouseSettings,
         logger: this.logger,
+        logLevel: config.logLevel,
+        keepAlive: config.keepAlive,
+        httpAgent: config.httpAgent,
       });
       this.reader = client;
       this.writer = client;
@@ -58,12 +70,18 @@ export class ClickHouse {
         url: config.readerUrl,
         clickhouseSettings: config.clickhouseSettings,
         logger: this.logger,
+        logLevel: config.logLevel,
+        keepAlive: config.keepAlive,
+        httpAgent: config.httpAgent,
       });
       this.writer = new ClickhouseClient({
         name: config.writerName ?? "clickhouse-writer",
         url: config.writerUrl,
         clickhouseSettings: config.clickhouseSettings,
         logger: this.logger,
+        logLevel: config.logLevel,
+        keepAlive: config.keepAlive,
+        httpAgent: config.httpAgent,
       });
 
       this._splitClients = true;