From fa0ce58e9c51b098d7db83eb5bbb11c051d7eadd Mon Sep 17 00:00:00 2001
From: Dotan Simha <dotansimha@gmail.com>
Date: Mon, 26 Feb 2024 08:47:28 +0200
Subject: [PATCH] Fix ClickHouse stats reporting (#4067)

---
 .../src/modules/operations/providers/clickhouse-client.ts    | 4 ++--
 .../services/api/src/modules/operations/providers/tokens.ts  | 2 +-
 packages/services/server/src/index.ts                        | 5 ++++-
 packages/services/usage-estimator/src/estimator.ts           | 2 +-
 packages/services/usage-estimator/src/index.ts               | 5 ++++-
 5 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/packages/services/api/src/modules/operations/providers/clickhouse-client.ts b/packages/services/api/src/modules/operations/providers/clickhouse-client.ts
index 8388d9e65..0aba7cfbe 100644
--- a/packages/services/api/src/modules/operations/providers/clickhouse-client.ts
+++ b/packages/services/api/src/modules/operations/providers/clickhouse-client.ts
@@ -18,7 +18,7 @@ function hashQuery(query: SqlStatement): string {
 export interface QueryResponse<T> {
   data: readonly T[];
   rows: number;
-  statistics: {
+  statistics?: {
     elapsed: number;
   };
 }
@@ -167,7 +167,7 @@ export class ClickHouse {
 
     this.config.onReadEnd?.(queryId, {
       totalSeconds: endedAt,
-      elapsedSeconds: response.statistics.elapsed,
+      elapsedSeconds: response.statistics?.elapsed,
     });
 
     return response;
diff --git a/packages/services/api/src/modules/operations/providers/tokens.ts b/packages/services/api/src/modules/operations/providers/tokens.ts
index 1a2fa04dc..6a9cab389 100644
--- a/packages/services/api/src/modules/operations/providers/tokens.ts
+++ b/packages/services/api/src/modules/operations/providers/tokens.ts
@@ -14,7 +14,7 @@ export interface ClickHouseConfig {
     label: string,
     timings: {
       totalSeconds: number;
-      elapsedSeconds: number;
+      elapsedSeconds?: number;
     },
   ) => void;
 }
diff --git a/packages/services/server/src/index.ts b/packages/services/server/src/index.ts
index e10eba4dc..0f7ad6e42 100644
--- a/packages/services/server/src/index.ts
+++ b/packages/services/server/src/index.ts
@@ -270,7 +270,10 @@ export async function main() {
         requestTimeout: env.clickhouse.requestTimeout,
         onReadEnd(query, timings) {
           clickHouseReadDuration.labels({ query }).observe(timings.totalSeconds);
-          clickHouseElapsedDuration.labels({ query }).observe(timings.elapsedSeconds);
+
+          if (timings.elapsedSeconds !== undefined) {
+            clickHouseElapsedDuration.labels({ query }).observe(timings.elapsedSeconds);
+          }
         },
       },
       cdn: env.cdn,
diff --git a/packages/services/usage-estimator/src/estimator.ts b/packages/services/usage-estimator/src/estimator.ts
index 3caa81ddb..48b0b101a 100644
--- a/packages/services/usage-estimator/src/estimator.ts
+++ b/packages/services/usage-estimator/src/estimator.ts
@@ -15,7 +15,7 @@ export function createEstimator(config: {
       label: string,
       timings: {
         totalSeconds: number;
-        elapsedSeconds: number;
+        elapsedSeconds?: number;
       },
     ) => void;
   };
diff --git a/packages/services/usage-estimator/src/index.ts b/packages/services/usage-estimator/src/index.ts
index 0388a0dc2..0386de091 100644
--- a/packages/services/usage-estimator/src/index.ts
+++ b/packages/services/usage-estimator/src/index.ts
@@ -46,7 +46,10 @@ async function main() {
         password: env.clickhouse.password,
         onReadEnd(query, timings) {
           clickHouseReadDuration.labels({ query }).observe(timings.totalSeconds);
-          clickHouseElapsedDuration.labels({ query }).observe(timings.elapsedSeconds);
+
+          if (timings.elapsedSeconds !== undefined) {
+            clickHouseElapsedDuration.labels({ query }).observe(timings.elapsedSeconds);
+          }
         },
       },
     });