From 8ff2e84ddbe12be797b2af223173edcc65a73599 Mon Sep 17 00:00:00 2001
From: blam <ben@blam.sh>
Date: Fri, 22 Sep 2023 16:10:59 +0200
Subject: [PATCH] feat: move over these metrics to open telemetry too

Signed-off-by: blam <ben@blam.sh>
---
 .../catalog-backend/src/database/metrics.ts   | 62 ++++++++++++++++---
 1 file changed, 55 insertions(+), 7 deletions(-)

diff --git a/plugins/catalog-backend/src/database/metrics.ts b/plugins/catalog-backend/src/database/metrics.ts
index b6c4e248dc..58ad91f93b 100644
--- a/plugins/catalog-backend/src/database/metrics.ts
+++ b/plugins/catalog-backend/src/database/metrics.ts
@@ -17,11 +17,14 @@
 import { Knex } from 'knex';
 import { createGaugeMetric } from '../util/metrics';
 import { DbRefreshStateRow, DbRelationsRow, DbLocationsRow } from './tables';
+import { metrics } from '@opentelemetry/api';
 
 export function initDatabaseMetrics(knex: Knex) {
+  const seenProm = new Set<string>();
   const seen = new Set<string>();
+  const meter = metrics.getMeter('default');
   return {
-    entities_count: createGaugeMetric({
+    entities_count_prom: createGaugeMetric({
       name: 'catalog_entities_count',
       help: 'Total amount of entities in the catalog',
       labelNames: ['kind'],
@@ -34,20 +37,20 @@ export function initDatabaseMetrics(knex: Knex) {
           .reduce((acc, e) => acc.set(e, (acc.get(e) || 0) + 1), new Map());
 
         results.forEach((value, key) => {
-          seen.add(key);
+          seenProm.add(key);
           this.set({ kind: key }, value);
         });
 
-        // Set all the entities that were not seen to 0 and delete them from the seen set.
-        seen.forEach(key => {
+        // Set all the entities that were not seenProm to 0 and delete them from the seenProm set.
+        seenProm.forEach(key => {
           if (!results.has(key)) {
             this.set({ kind: key }, 0);
-            seen.delete(key);
+            seenProm.delete(key);
           }
         });
       },
     }),
-    registered_locations: createGaugeMetric({
+    registered_locations_prom: createGaugeMetric({
       name: 'catalog_registered_locations_count',
       help: 'Total amount of registered locations in the catalog',
       async collect() {
@@ -57,7 +60,7 @@ export function initDatabaseMetrics(knex: Knex) {
         this.set(Number(total[0].count));
       },
     }),
-    relations: createGaugeMetric({
+    relations_prom: createGaugeMetric({
       name: 'catalog_relations_count',
       help: 'Total amount of relations between entities',
       async collect() {
@@ -67,5 +70,50 @@ export function initDatabaseMetrics(knex: Knex) {
         this.set(Number(total[0].count));
       },
     }),
+    entities_count: meter
+      .createObservableGauge('catalog_entities_count', {
+        description: 'Total amount of entities in the catalog',
+      })
+      .addCallback(async gauge => {
+        const result = await knex<DbRefreshStateRow>('refresh_state').select(
+          'entity_ref',
+        );
+        const results = result
+          .map(row => row.entity_ref.split(':')[0])
+          .reduce((acc, e) => acc.set(e, (acc.get(e) || 0) + 1), new Map());
+
+        results.forEach((value, key) => {
+          seen.add(key);
+          gauge.observe(value, { kind: key });
+        });
+
+        // Set all the entities that were not seen to 0 and delete them from the seen set.
+        seen.forEach(key => {
+          if (!results.has(key)) {
+            gauge.observe(0, { kind: key });
+            seen.delete(key);
+          }
+        });
+      }),
+    registered_locations: meter
+      .createObservableGauge('catalog_registered_locations_count', {
+        description: 'Total amount of registered locations in the catalog',
+      })
+      .addCallback(async gauge => {
+        const total = await knex<DbLocationsRow>('locations').count({
+          count: '*',
+        });
+        gauge.observe(Number(total[0].count));
+      }),
+    relations: meter
+      .createObservableGauge('catalog_relations_count', {
+        description: 'Total amount of relations between entities',
+      })
+      .addCallback(async gauge => {
+        const total = await knex<DbRelationsRow>('relations').count({
+          count: '*',
+        });
+        gauge.observe(Number(total[0].count));
+      }),
   };
 }