Some fixes and cleanups

Author: ericallam

.configs/prometheus.yml

@@ -48,6 +48,10 @@ export class ClickhouseClient implements ClickhouseReader, ClickhouseWriter {
     this.tracer = config.tracer ?? trace.getTracer("@internal/clickhouse");
   }
 
+  public async close() {
+    await this.client.close();
+  }
+
   public query<TIn extends z.ZodSchema<any>, TOut extends z.ZodSchema<any>>(req: {
     /**
      * The name of the operation.

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

@@ -6,6 +6,10 @@ import { z } from "zod";
 import { ClickHouseSettings, InsertResult } from "@clickhouse/client";
 
 export class NoopClient implements ClickhouseReader, ClickhouseWriter {
+  public async close() {
+    return;
+  }
+
   public query<TIn extends z.ZodSchema<any>, TOut extends z.ZodSchema<any>>(req: {
     query: string;
     params?: TIn;

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

@@ -41,6 +41,8 @@ export interface ClickhouseReader {
      */
     settings?: ClickHouseSettings;
   }): ClickhouseQueryFunction<z.input<TIn>, z.output<TOut>>;
+
+  close(): Promise<void>;
 }
 
 export type ClickhouseInsertFunction<TInput> = (
@@ -58,4 +60,6 @@ export interface ClickhouseWriter {
     schema: TSchema;
     settings?: ClickHouseSettings;
   }): ClickhouseInsertFunction<z.input<TSchema>>;
+
+  close(): Promise<void>;
 }

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

@@ -31,6 +31,7 @@ export class ClickHouse {
   public readonly reader: ClickhouseReader;
   public readonly writer: ClickhouseWriter;
   private readonly logger: Logger;
+  private _splitClients: boolean;
 
   constructor(config: ClickHouseConfig) {
     this.logger = config.logger ?? new Logger("ClickHouse", "debug");
@@ -49,6 +50,8 @@ export class ClickHouse {
       });
       this.reader = client;
       this.writer = client;
+
+      this._splitClients = false;
     } else if (config.writerUrl && config.readerUrl) {
       this.reader = new ClickhouseClient({
         name: config.readerName ?? "clickhouse-reader",
@@ -62,9 +65,13 @@ export class ClickHouse {
         clickhouseSettings: config.clickhouseSettings,
         logger: this.logger,
       });
+
+      this._splitClients = true;
     } else {
       this.reader = new NoopClient();
       this.writer = new NoopClient();
+
+      this._splitClients = true;
     }
   }
 
@@ -87,6 +94,14 @@ export class ClickHouse {
     });
   }
 
+  async close() {
+    if (this._splitClients) {
+      await Promise.all([this.reader.close(), this.writer.close()]);
+    } else {
+      await this.reader.close();
+    }
+  }
+
   get taskRuns() {
     return {
       insert: insertTaskRuns(this.writer),

internal-packages/clickhouse/src/index.ts

@@ -1,4 +1,3 @@
 export * from "./client.js";
 export * from "./errors.js";
-export * from "./stream.js";
 export type * from "./pgoutput.js";