From 743b06ff2398c01511a4a7eb48364700bcf6d78d Mon Sep 17 00:00:00 2001
From: Sam DeHaan <sam.dehaan@grafana.com>
Date: Thu, 24 Oct 2024 17:50:57 +0200
Subject: [PATCH] Ensure database connections are always closed

Signed-off-by: Sam DeHaan <sam.dehaan@grafana.com>
---
 cmd/postgres_exporter/server.go | 1 +
 collector/collector.go          | 6 +++++-
 collector/probe.go              | 2 +-
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/cmd/postgres_exporter/server.go b/cmd/postgres_exporter/server.go
index bcfee6812..3f2074490 100644
--- a/cmd/postgres_exporter/server.go
+++ b/cmd/postgres_exporter/server.go
@@ -169,6 +169,7 @@ func (s *Servers) GetServer(dsn string) (*Server, error) {
 			s.servers[dsn] = server
 		}
 		if err = server.Ping(); err != nil {
+			server.Close()
 			delete(s.servers, dsn)
 			time.Sleep(time.Duration(errCount) * time.Second)
 			continue
diff --git a/collector/collector.go b/collector/collector.go
index 121129871..7b7a39d28 100644
--- a/collector/collector.go
+++ b/collector/collector.go
@@ -172,11 +172,11 @@ func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
 
 	// Set up the database connection for the collector.
 	err := inst.setup()
+	defer inst.Close()
 	if err != nil {
 		level.Error(p.logger).Log("msg", "Error opening connection to database", "err", err)
 		return
 	}
-	defer inst.Close()
 
 	wg := sync.WaitGroup{}
 	wg.Add(len(p.Collectors))
@@ -189,6 +189,10 @@ func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
 	wg.Wait()
 }
 
+func (p *PostgresCollector) Close() error {
+	return p.instance.Close()
+}
+
 func execute(ctx context.Context, name string, c Collector, instance *instance, ch chan<- prometheus.Metric, logger log.Logger) {
 	begin := time.Now()
 	err := c.Update(ctx, instance, ch)
diff --git a/collector/probe.go b/collector/probe.go
index 4c0f0419b..5401f687d 100644
--- a/collector/probe.go
+++ b/collector/probe.go
@@ -77,11 +77,11 @@ func (pc *ProbeCollector) Describe(ch chan<- *prometheus.Desc) {
 func (pc *ProbeCollector) Collect(ch chan<- prometheus.Metric) {
 	// Set up the database connection for the collector.
 	err := pc.instance.setup()
+	defer pc.instance.Close()
 	if err != nil {
 		level.Error(pc.logger).Log("msg", "Error opening connection to database", "err", err)
 		return
 	}
-	defer pc.instance.Close()
 
 	wg := sync.WaitGroup{}
 	wg.Add(len(pc.collectors))