From ea98ac042e0fa186f912e2e0796f370b55137cc2 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Wed, 17 Nov 2021 15:40:02 +0100
Subject: [PATCH] Add Prometheus instrumentation for the database

---
 cmd/soju/main.go | 4 ++++
 db.go            | 7 +++++++
 db_postgres.go   | 6 ++++++
 db_sqlite.go     | 6 ++++++
 4 files changed, 23 insertions(+)

diff --git a/cmd/soju/main.go b/cmd/soju/main.go
index bd07336..214c0cc 100644
--- a/cmd/soju/main.go
+++ b/cmd/soju/main.go
@@ -293,6 +293,10 @@ func main() {
 		log.Printf("server listening on %q", listen)
 	}
 
+	if db, ok := db.(soju.MetricsCollectorDatabase); ok && srv.MetricsRegistry != nil {
+		srv.MetricsRegistry.MustRegister(db.MetricsCollector())
+	}
+
 	sigCh := make(chan os.Signal, 1)
 	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
 
diff --git a/db.go b/db.go
index c2c845d..41dd42e 100644
--- a/db.go
+++ b/db.go
@@ -6,6 +6,8 @@ import (
 	"net/url"
 	"strings"
 	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
 )
 
 type Database interface {
@@ -28,6 +30,11 @@ type Database interface {
 	StoreClientDeliveryReceipts(ctx context.Context, networkID int64, client string, receipts []DeliveryReceipt) error
 }
 
+type MetricsCollectorDatabase interface {
+	Database
+	MetricsCollector() prometheus.Collector
+}
+
 func OpenDB(driver, source string) (Database, error) {
 	switch driver {
 	case "sqlite3":
diff --git a/db_postgres.go b/db_postgres.go
index 7db1edc..5ede40e 100644
--- a/db_postgres.go
+++ b/db_postgres.go
@@ -10,6 +10,8 @@ import (
 	"time"
 
 	_ "github.com/lib/pq"
+	"github.com/prometheus/client_golang/prometheus"
+	promcollectors "github.com/prometheus/client_golang/prometheus/collectors"
 )
 
 const postgresQueryTimeout = 5 * time.Second
@@ -148,6 +150,10 @@ func (db *PostgresDB) Close() error {
 	return db.db.Close()
 }
 
+func (db *PostgresDB) MetricsCollector() prometheus.Collector {
+	return promcollectors.NewDBStatsCollector(db.db, "main")
+}
+
 func (db *PostgresDB) Stats(ctx context.Context) (*DatabaseStats, error) {
 	ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
 	defer cancel()
diff --git a/db_sqlite.go b/db_sqlite.go
index 1616c1e..2f4d99f 100644
--- a/db_sqlite.go
+++ b/db_sqlite.go
@@ -10,6 +10,8 @@ import (
 	"time"
 
 	_ "github.com/mattn/go-sqlite3"
+	"github.com/prometheus/client_golang/prometheus"
+	promcollectors "github.com/prometheus/client_golang/prometheus/collectors"
 )
 
 const sqliteQueryTimeout = 5 * time.Second
@@ -238,6 +240,10 @@ func (db *SqliteDB) upgrade() error {
 	return tx.Commit()
 }
 
+func (db *SqliteDB) MetricsCollector() prometheus.Collector {
+	return promcollectors.NewDBStatsCollector(db.db, "main")
+}
+
 func (db *SqliteDB) Stats(ctx context.Context) (*DatabaseStats, error) {
 	db.lock.RLock()
 	defer db.lock.RUnlock()