Merge pull request #785 from sysadmind/collector-instance

Add the instance struct to handle connections

Author: SuperQ

collector/collector.go

@@ -15,7 +15,6 @@ package collector
 
 import (
 "context"
-"database/sql"
 "errors"
 "fmt"
 "sync"
@@ -59,7 +58,7 @@ var (
 )
 
 type Collector interface {
-Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error
+Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error
 }
 
 type collectorConfig struct {
@@ -92,7 +91,7 @@ type PostgresCollector struct {
 Collectors map[string]Collector
 logger log.Logger
 
-db *sql.DB
+instance *instance
 }
 
 type Option func(*PostgresCollector) error
@@ -149,14 +148,11 @@ func NewPostgresCollector(logger log.Logger, excludeDatabases []string, dsn stri
 return nil, errors.New("empty dsn")
 }
 
-db, err := sql.Open("postgres", dsn)
+instance, err := newInstance(dsn)
 if err != nil {
 return nil, err
 }
-db.SetMaxOpenConns(1)
-db.SetMaxIdleConns(1)
-
-p.db = db
+p.instance = instance
 
 return p, nil
 }
@@ -174,16 +170,16 @@ func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
 wg.Add(len(p.Collectors))
 for name, c := range p.Collectors {
 go func(name string, c Collector) {
-execute(ctx, name, c, p.db, ch, p.logger)
+execute(ctx, name, c, p.instance, ch, p.logger)
 wg.Done()
 }(name, c)
 }
 wg.Wait()
 }
 
-func execute(ctx context.Context, name string, c Collector, db *sql.DB, ch chan<- prometheus.Metric, logger log.Logger) {
+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, db, ch)
+err := c.Update(ctx, instance, ch)
 duration := time.Since(begin)
 var success float64
 

collector/instance.go

@@ -0,0 +1,85 @@
+// Copyright 2023 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package collector
+
+import (
+"database/sql"
+"fmt"
+"regexp"
+
+"github.com/blang/semver/v4"
+)
+
+type instance struct {
+db *sql.DB
+version semver.Version
+}
+
+func newInstance(dsn string) (*instance, error) {
+i := &instance{}
+db, err := sql.Open("postgres", dsn)
+if err != nil {
+return nil, err
+}
+db.SetMaxOpenConns(1)
+db.SetMaxIdleConns(1)
+i.db = db
+
+version, err := queryVersion(db)
+if err != nil {
+db.Close()
+return nil, err
+}
+
+i.version = version
+
+return i, nil
+}
+
+func (i *instance) getDB() *sql.DB {
+return i.db
+}
+
+func (i *instance) Close() error {
+return i.db.Close()
+}
+
+// Regex used to get the "short-version" from the postgres version field.
+// The result of SELECT version() is something like "PostgreSQL 9.6.2 on x86_64-pc-linux-gnu, compiled by gcc (GCC) 6.2.1 20160830, 64-bit"
+var versionRegex = regexp.MustCompile(`^\w+ ((\d+)(\.\d+)?(\.\d+)?)`)
+var serverVersionRegex = regexp.MustCompile(`^((\d+)(\.\d+)?(\.\d+)?)`)
+
+func queryVersion(db *sql.DB) (semver.Version, error) {
+var version string
+err := db.QueryRow("SELECT version();").Scan(&version)
+if err != nil {
+return semver.Version{}, err
+}
+submatches := versionRegex.FindStringSubmatch(version)
+if len(submatches) > 1 {
+return semver.ParseTolerant(submatches[1])
+}
+
+// We could also try to parse the version from the server_version field.
+// This is of the format 13.3 (Debian 13.3-1.pgdg100+1)
+err = db.QueryRow("SHOW server_version;").Scan(&version)
+if err != nil {
+return semver.Version{}, err
+}
+submatches = serverVersionRegex.FindStringSubmatch(version)
+if len(submatches) > 1 {
+return semver.ParseTolerant(submatches[1])
+}
+return semver.Version{}, fmt.Errorf("could not parse version from %q", version)
+}

collector/pg_database.go

@@ -15,7 +15,6 @@ package collector
 
 import (
 "context"
-"database/sql"
 
 "github.com/go-kit/log"
 "github.com/prometheus/client_golang/prometheus"
@@ -66,7 +65,8 @@ var (
 // each database individually. This is because we can't filter the
 // list of databases in the query because the list of excluded
 // databases is dynamic.
-func (c PGDatabaseCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
+func (c PGDatabaseCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
+db := instance.getDB()
 // Query the list of databases
 rows, err := db.QueryContext(ctx,
 pgDatabaseQuery,

collector/pg_database_test.go

@@ -29,6 +29,8 @@ func TestPGDatabaseCollector(t *testing.T) {
 }
 defer db.Close()
 
+inst := &instance{db: db}
+
 mock.ExpectQuery(sanitizeQuery(pgDatabaseQuery)).WillReturnRows(sqlmock.NewRows([]string{"datname"}).
 AddRow("postgres"))
 
@@ -39,7 +41,7 @@ func TestPGDatabaseCollector(t *testing.T) {
 go func() {
 defer close(ch)
 c := PGDatabaseCollector{}
-if err := c.Update(context.Background(), db, ch); err != nil {
+if err := c.Update(context.Background(), inst, ch); err != nil {
 t.Errorf("Error calling PGDatabaseCollector.Update: %s", err)
 }
 }()

collector/pg_postmaster.go

@@ -15,7 +15,6 @@ package collector
 
 import (
 "context"
-"database/sql"
 
 "github.com/prometheus/client_golang/prometheus"
 )
@@ -47,7 +46,8 @@ var (
 pgPostmasterQuery = "SELECT pg_postmaster_start_time from pg_postmaster_start_time();"
 )
 
-func (c *PGPostmasterCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
+func (c *PGPostmasterCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
+db := instance.getDB()
 row := db.QueryRowContext(ctx,
 pgPostmasterQuery)
 

collector/pg_postmaster_test.go

@@ -29,6 +29,8 @@ func TestPgPostmasterCollector(t *testing.T) {
 }
 defer db.Close()
 
+inst := &instance{db: db}
+
 mock.ExpectQuery(sanitizeQuery(pgPostmasterQuery)).WillReturnRows(sqlmock.NewRows([]string{"pg_postmaster_start_time"}).
 AddRow(1685739904))
 
@@ -37,7 +39,7 @@ func TestPgPostmasterCollector(t *testing.T) {
 defer close(ch)
 c := PGPostmasterCollector{}
 
-if err := c.Update(context.Background(), db, ch); err != nil {
+if err := c.Update(context.Background(), inst, ch); err != nil {
 t.Errorf("Error calling PGPostmasterCollector.Update: %s", err)
 }
 }()

collector/pg_process_idle.go

@@ -15,7 +15,6 @@ package collector
 
 import (
 "context"
-"database/sql"
 
 "github.com/go-kit/log"
 "github.com/prometheus/client_golang/prometheus"
@@ -42,7 +41,8 @@ var pgProcessIdleSeconds = prometheus.NewDesc(
 prometheus.Labels{},
 )
 
-func (PGProcessIdleCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
+func (PGProcessIdleCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
+db := instance.getDB()
 row := db.QueryRowContext(ctx,
 `WITH
 metrics AS (

collector/pg_replication_slot.go

@@ -15,7 +15,6 @@ package collector
 
 import (
 "context"
-"database/sql"
 
 "github.com/go-kit/log"
 "github.com/prometheus/client_golang/prometheus"
@@ -73,7 +72,8 @@ var (
 pg_replication_slots;`
 )
 
-func (PGReplicationSlotCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
+func (PGReplicationSlotCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
+db := instance.getDB()
 rows, err := db.QueryContext(ctx,
 pgReplicationSlotQuery)
 if err != nil {

collector/pg_replication_slot_test.go

@@ -29,6 +29,8 @@ func TestPgReplicationSlotCollectorActive(t *testing.T) {
 }
 defer db.Close()
 
+inst := &instance{db: db}
+
 columns := []string{"slot_name", "current_wal_lsn", "confirmed_flush_lsn", "active"}
 rows := sqlmock.NewRows(columns).
 AddRow("test_slot", 5, 3, true)
@@ -39,7 +41,7 @@ func TestPgReplicationSlotCollectorActive(t *testing.T) {
 defer close(ch)
 c := PGReplicationSlotCollector{}
 
-if err := c.Update(context.Background(), db, ch); err != nil {
+if err := c.Update(context.Background(), inst, ch); err != nil {
 t.Errorf("Error calling PGPostmasterCollector.Update: %s", err)
 }
 }()
@@ -68,6 +70,8 @@ func TestPgReplicationSlotCollectorInActive(t *testing.T) {
 }
 defer db.Close()
 
+inst := &instance{db: db}
+
 columns := []string{"slot_name", "current_wal_lsn", "confirmed_flush_lsn", "active"}
 rows := sqlmock.NewRows(columns).
 AddRow("test_slot", 6, 12, false)
@@ -78,7 +82,7 @@ func TestPgReplicationSlotCollectorInActive(t *testing.T) {
 defer close(ch)
 c := PGReplicationSlotCollector{}
 
-if err := c.Update(context.Background(), db, ch); err != nil {
+if err := c.Update(context.Background(), inst, ch); err != nil {
 t.Errorf("Error calling PGReplicationSlotCollector.Update: %s", err)
 }
 }()

collector/pg_stat_bgwriter.go

@@ -15,7 +15,6 @@ package collector
 
 import (
 "context"
-"database/sql"
 "time"
 
 "github.com/prometheus/client_golang/prometheus"
@@ -117,7 +116,8 @@ var (
 FROM pg_stat_bgwriter;`
 )
 
-func (PGStatBGWriterCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
+func (PGStatBGWriterCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
+db := instance.getDB()
 row := db.QueryRowContext(ctx,
 statBGWriterQuery)