From 5f781b4c8f60100910b4f29c08cae26d6e5b64c9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nuno=20G=C3=B3is?= <github@nunogois.com>
Date: Fri, 16 Feb 2024 15:11:29 +0000
Subject: [PATCH] refactor: better prom metric helper types (#6261)

Improves typing in our Prometheus metric helpers.
---
 src/lib/util/metrics/createCounter.ts | 16 +++++++++++++---
 src/lib/util/metrics/createGauge.ts   | 16 +++++++++++++---
 src/lib/util/metrics/createSummary.ts | 15 ++++++++++++---
 3 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/src/lib/util/metrics/createCounter.ts b/src/lib/util/metrics/createCounter.ts
index 16084c34dd..0de98867d9 100644
--- a/src/lib/util/metrics/createCounter.ts
+++ b/src/lib/util/metrics/createCounter.ts
@@ -1,4 +1,14 @@
-import { Counter, CounterConfiguration } from 'prom-client';
+import { Counter as PromCounter, CounterConfiguration } from 'prom-client';
+
+/**
+ * A wrapped instance of prom-client's Counter, overriding some of its methods for enhanced functionality and type-safety.
+ */
+export type Counter<T extends string = string> = {
+    counter: PromCounter<T>;
+    labels: (labels: Record<T, string | number>) => PromCounter.Internal;
+    inc: (value?: number | undefined) => void;
+    increment: (labels: Record<T, string | number>, value?: number) => void;
+};
 
 /**
  * Creates a wrapped instance of prom-client's Counter, overriding some of its methods for enhanced functionality and type-safety.
@@ -9,11 +19,11 @@ import { Counter, CounterConfiguration } from 'prom-client';
  */
 export const createCounter = <T extends string>(
     options: CounterConfiguration<T>,
-) => {
+): Counter<T> => {
     /**
      * The underlying instance of prom-client's Counter.
      */
-    const counter = new Counter(options);
+    const counter = new PromCounter<T>(options);
 
     /**
      * Applies given labels to the counter. Labels are key-value pairs.
diff --git a/src/lib/util/metrics/createGauge.ts b/src/lib/util/metrics/createGauge.ts
index ef0b078366..383fbab462 100644
--- a/src/lib/util/metrics/createGauge.ts
+++ b/src/lib/util/metrics/createGauge.ts
@@ -1,4 +1,14 @@
-import { Gauge, GaugeConfiguration } from 'prom-client';
+import { Gauge as PromGauge, GaugeConfiguration } from 'prom-client';
+
+/**
+ * A wrapped instance of prom-client's Gauge, overriding some of its methods for enhanced functionality and type-safety.
+ */
+export type Gauge<T extends string = string> = {
+    gauge: PromGauge<T>;
+    labels: (labels: Record<T, string | number>) => PromGauge.Internal<T>;
+    reset: () => void;
+    set: (value: number) => void;
+};
 
 /**
  * Creates a wrapped instance of prom-client's Gauge, overriding some of its methods for enhanced functionality and type-safety.
@@ -9,11 +19,11 @@ import { Gauge, GaugeConfiguration } from 'prom-client';
  */
 export const createGauge = <T extends string>(
     options: GaugeConfiguration<T>,
-) => {
+): Gauge<T> => {
     /**
      * The underlying instance of prom-client's Gauge.
      */
-    const gauge = new Gauge(options);
+    const gauge = new PromGauge(options);
 
     /**
      * Applies given labels to the gauge. Labels are key-value pairs.
diff --git a/src/lib/util/metrics/createSummary.ts b/src/lib/util/metrics/createSummary.ts
index bdd4cc47b9..cab3706a83 100644
--- a/src/lib/util/metrics/createSummary.ts
+++ b/src/lib/util/metrics/createSummary.ts
@@ -1,4 +1,13 @@
-import { Summary, SummaryConfiguration } from 'prom-client';
+import { Summary as PromSummary, SummaryConfiguration } from 'prom-client';
+
+/**
+ * A wrapped instance of prom-client's Summary, overriding some of its methods for enhanced functionality and type-safety.
+ */
+export type Summary<T extends string = string> = {
+    summary: PromSummary<T>;
+    labels: (labels: Record<T, string | number>) => PromSummary.Internal<T>;
+    observe: (value: number) => void;
+};
 
 /**
  * Creates a wrapped instance of prom-client's Summary, overriding some of its methods for enhanced functionality and type-safety.
@@ -9,11 +18,11 @@ import { Summary, SummaryConfiguration } from 'prom-client';
  */
 export const createSummary = <T extends string>(
     options: SummaryConfiguration<T>,
-) => {
+): Summary<T> => {
     /**
      * The underlying instance of prom-client's Summary.
      */
-    const summary = new Summary(options);
+    const summary = new PromSummary(options);
 
     /**
      * Applies given labels to the summary. Labels are key-value pairs.